import time from typing import Optional from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, delete from app.core.biz_exception import NotFoundError from app.models.account import Account from app.schemas.account import AccountCreate, LockRequest class AccountService: @staticmethod async def add_account(db: AsyncSession, payload: AccountCreate): """ 添加新账号到数据库 """ # 检查是否存在 stmt = select(Account).where( Account.pool_name == payload.pool_name, Account.username == payload.username, ) obj = (await db.execute(stmt)).scalar_one_or_none() if obj: # 如果存在,更新密码,并重置为 active obj.password = payload.password obj.status = "active" if payload.extra_data: obj.extra_data = payload.extra_data await db.commit() await db.refresh(obj) return obj else: new_acc = Account( pool_name=payload.pool_name, username=payload.username, password=payload.password, extra_data=payload.extra_data ) db.add(new_acc) await db.commit() await db.refresh(new_acc) return new_acc @staticmethod async def get_next_account(db: AsyncSession, pool_name: str, lock_duration: float) -> Account: now = time.time() stmt = select(Account).where( Account.pool_name == pool_name, Account.status == 'active', Account.lock_until < now ).order_by( Account.lock_until.asc() ).limit(1).with_for_update() obj = (await db.execute(stmt)).scalar_one_or_none() if not obj: raise NotFoundError('Account not found') new_lock_time = now + lock_duration obj.lock_until = new_lock_time await db.commit() await db.refresh(obj) return obj @staticmethod async def manual_lock(db: AsyncSession, payload: LockRequest): stmt = select(Account).where( Account.pool_name == payload.pool_name, Account.username == payload.username, ) obj = (await db.execute(stmt)).scalar_one_or_none() if not obj: raise NotFoundError('Account not found') obj.lock_until = time.time() + payload.duration await db.commit() @staticmethod async def disable_account(db: AsyncSession, payload: LockRequest): stmt = select(Account).where( Account.pool_name == payload.pool_name, Account.username == payload.username, ) obj = (await db.execute(stmt)).scalar_one_or_none() if not obj: raise NotFoundError('Account not found') obj.status = "disabled" await db.commit()