69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""账号额度与 quota_disabled 同步。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from models.models import Account, User
|
|
from .account_limits import account_limit_for_user
|
|
|
|
QUOTA_DISABLED_MESSAGE = "账号额度不足，该账号已被系统停用"
|
|
|
|
|
|
async def sync_user_account_quota(
|
|
db: AsyncSession,
|
|
user: User,
|
|
*,
|
|
stop_worker: Callable[[int], Awaitable[None]] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""按 max_accounts 同步停用状态：保留最早创建的账号，超额账号禁用并停止托管。"""
|
|
limit = account_limit_for_user(user)
|
|
|
|
result = await db.execute(
|
|
select(Account)
|
|
.where(Account.owner_id == user.id)
|
|
.order_by(Account.created_at.asc(), Account.id.asc())
|
|
)
|
|
accounts = list(result.scalars().all())
|
|
|
|
disabled_ids: list[int] = []
|
|
enabled_ids: list[int] = []
|
|
|
|
if limit is None:
|
|
for acc in accounts:
|
|
if acc.quota_disabled:
|
|
acc.quota_disabled = False
|
|
if (acc.error_message or "").strip() == QUOTA_DISABLED_MESSAGE:
|
|
acc.error_message = None
|
|
enabled_ids.append(acc.id)
|
|
return {"disabled_ids": disabled_ids, "enabled_ids": enabled_ids, "limit": None}
|
|
|
|
for idx, acc in enumerate(accounts):
|
|
if idx < limit:
|
|
if acc.quota_disabled:
|
|
acc.quota_disabled = False
|
|
if (acc.error_message or "").strip() == QUOTA_DISABLED_MESSAGE:
|
|
acc.error_message = None
|
|
enabled_ids.append(acc.id)
|
|
else:
|
|
if not acc.quota_disabled:
|
|
acc.quota_disabled = True
|
|
acc.error_message = QUOTA_DISABLED_MESSAGE
|
|
acc.status = "offline"
|
|
acc.qr_code_base64 = None
|
|
disabled_ids.append(acc.id)
|
|
if stop_worker:
|
|
await stop_worker(acc.id)
|
|
|
|
return {"disabled_ids": disabled_ids, "enabled_ids": enabled_ids, "limit": limit}
|
|
|
|
|
|
async def default_stop_worker(account_id: int) -> None:
|
|
from main import manager
|
|
|
|
await manager.stop_worker(account_id)
|