133 lines
3.6 KiB
Python
133 lines
3.6 KiB
Python
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
|