webapp/app/services/account_service.py

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()