from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.core.roles import ROLE_AGENT from app.models.card_key import CardKey from app.models.user import User from app.schemas.oem_card_key import OemCardKeyBatchAssign, OemCardKeyUpdate DEFAULT_PAGE_SIZE = 20 MAX_PAGE_SIZE = 200 class OemCardKeyError(Exception): def __init__(self, message: str) -> None: self.message = message super().__init__(message) def _list_card_keys_filters( *, oem_id: int, status: str | None = None, username: str | None = None, agent_id: int | None = None, ) -> list: filters = [CardKey.oem_id == oem_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()}%")) if agent_id is not None: filters.append(CardKey.agent_id == agent_id) return filters async def list_card_keys( db: AsyncSession, *, oem_id: int, page: int = 1, page_size: int = DEFAULT_PAGE_SIZE, status: str | None = None, username: str | None = None, agent_id: int | 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( oem_id=oem_id, status=status, username=username, agent_id=agent_id, ) 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_oem(db: AsyncSession, card_id: int, oem_id: int) -> CardKey: card = await db.get(CardKey, card_id) if card is None or card.oem_id != oem_id: raise OemCardKeyError("卡密不存在") return card async def _validate_agent_under_oem( db: AsyncSession, *, oem_id: int, agent_id: int ) -> User: agent = await db.get(User, agent_id) if agent is None or agent.role_id != ROLE_AGENT: raise OemCardKeyError("所选代理不存在") if agent.oem_id != oem_id: raise OemCardKeyError("所选代理不属于当前 OEM") return agent async def update_card_key( db: AsyncSession, card_id: int, body: OemCardKeyUpdate, *, oem_id: int, ) -> CardKey: card = await _get_card_under_oem(db, card_id, oem_id) if "oem_remark" in body.model_fields_set: card.oem_remark = body.oem_remark await db.flush() await db.refresh(card) return card async def batch_assign_to_agent( db: AsyncSession, body: OemCardKeyBatchAssign, *, oem_id: int, ) -> int: await _validate_agent_under_oem(db, oem_id=oem_id, agent_id=body.agent_id) result = await db.scalars( select(CardKey).where(CardKey.id.in_(body.card_ids)) ) cards = list(result.all()) if len(cards) != len(set(body.card_ids)): raise OemCardKeyError("部分卡密不存在") updated = 0 for card in cards: if card.oem_id != oem_id: raise OemCardKeyError("部分卡密不属于当前 OEM") if card.activated_at is not None: raise OemCardKeyError("已激活卡密不能分配") card.agent_id = body.agent_id updated += 1 await db.flush() return updated