ShiErFeng/shierfeng-fastapi-backend/module_business/service/test_batch_purge_service.py

import hashlib
import hmac
import json
from datetime import datetime, timedelta
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from common.vo import CrudResponseModel
from config.database import Base
from config.env import TestBatchConfig
from exceptions.exception import PermissionException, ServiceException, ServiceWarning
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_business.dao.test_batch_dao import TestBatchDao
from module_business.dao.test_batch_purge_dao import TestBatchPurgeDao
from module_business.domain.enums import TestBatchStatus
from module_business.domain.test_batch_graph import (
    BATCH_LABELS,
    BATCH_MODELS,
    BatchSelection,
    derive_batch_delete_order,
    issue_preview_token,
    selection_hash,
)
from module_business.entity.vo.test_batch_vo import (
    BatchImpactModel,
    PurgeExecuteRequestModel,
    PurgePreviewModel,
)
from module_business.service.test_batch_access_service import TestBatchAccessService


class TestBatchPurgeService:
    """预览并原子清理已关闭的测试批次。"""

    @staticmethod
    def _assert_admin(current_user: CurrentUserModel) -> None:
        TestBatchAccessService.assert_super_admin(current_user)

    @staticmethod
    def _user_identity(current_user: CurrentUserModel) -> tuple[int, str]:
        user = current_user.user
        if user is None or user.user_id is None or user.user_name is None:
            raise PermissionException('', '当前用户信息不可用')
        return user.user_id, user.user_name

    @staticmethod
    def _assert_closed(batch: Any | None) -> None:
        if batch is None:
            raise ServiceWarning('', '测试批次不存在')
        if batch.status == TestBatchStatus.PURGED.value:
            raise ServiceWarning('', '测试批次已清理')
        if batch.status == TestBatchStatus.ACTIVE.value:
            raise ServiceWarning('', '测试批次必须先关闭')
        if batch.status != TestBatchStatus.CLOSED.value:
            raise ServiceWarning('', '测试批次状态异常')

    @staticmethod
    def _assert_token(batch: object, raw_token: str, now: datetime) -> None:
        supplied_hash = hashlib.sha256(raw_token.encode()).hexdigest()
        stored_hash = getattr(batch, 'preview_token_hash', None)
        comparable_hash = stored_hash if isinstance(stored_hash, str) else '0' * 64
        if not hmac.compare_digest(comparable_hash, supplied_hash):
            raise ServiceWarning('', '预览令牌无效')
        if getattr(batch, 'preview_action', None) != 'PURGE':
            raise ServiceWarning('', '预览令牌无效')
        expires_at = getattr(batch, 'preview_expires_at', None)
        if expires_at is None or expires_at <= now:
            raise ServiceWarning('', '预览令牌已过期')

    @staticmethod
    def _counts(selection: BatchSelection) -> dict[str, int]:
        return {table_name: len(selection[table_name]) for table_name in BATCH_MODELS}

    @staticmethod
    def _blocking_reason(total_count: int, limit: int) -> str | None:
        if limit <= 0:
            return '尚未配置在线测试批次清理上限'
        if total_count > limit:
            return '批次数据量超过在线清理上限，请按运维流程处理'
        return None

    @classmethod
    async def preview(
        cls,
        db: AsyncSession,
        batch_id: int,
        current_user: CurrentUserModel,
    ) -> PurgePreviewModel:
        cls._assert_admin(current_user)
        _, user_name = cls._user_identity(current_user)
        try:
            batch = await TestBatchDao.get_batch_for_update(db, batch_id)
            cls._assert_closed(batch)
            selection = await TestBatchPurgeDao.load_selection(db, batch_id, for_update=False)
            display_keys = await TestBatchPurgeDao.list_display_keys(db, selection)
            counts = cls._counts(selection)
            total_count = sum(counts.values())
            limit = TestBatchConfig.test_batch_online_purge_max_rows
            blocking_reason = cls._blocking_reason(total_count, limit)

            now = datetime.now()
            expires_at = now + timedelta(seconds=TestBatchConfig.test_batch_preview_ttl_seconds)
            raw_token, token_hash = issue_preview_token()
            await TestBatchDao.store_preview(
                db,
                batch,
                action='PURGE',
                token_hash=token_hash,
                scope_hash=selection_hash(selection),
                expires_at=expires_at,
                update_by=user_name,
                update_time=now,
            )
            impacts = [
                BatchImpactModel(
                    table_name=table_name,
                    label=BATCH_LABELS[table_name],
                    count=counts[table_name],
                    keys=display_keys[table_name],
                )
                for table_name in BATCH_MODELS
            ]
            response = PurgePreviewModel(
                batch_id=batch_id,
                batch_name=batch.batch_name,
                preview_token=raw_token,
                expires_at=expires_at,
                total_count=total_count,
                impacts=impacts,
                online_purge_allowed=blocking_reason is None,
                online_purge_limit=limit,
                blocking_reason=blocking_reason,
            )
            await db.commit()
            return response
        except Exception:
            await db.rollback()
            raise

    @classmethod
    async def purge(
        cls,
        db: AsyncSession,
        batch_id: int,
        command: PurgeExecuteRequestModel,
        current_user: CurrentUserModel,
    ) -> CrudResponseModel:
        cls._assert_admin(current_user)
        user_id, user_name = cls._user_identity(current_user)
        try:
            batch = await TestBatchDao.get_batch_for_update(db, batch_id)
            if batch is not None and batch.status == TestBatchStatus.PURGED.value:
                response = CrudResponseModel(is_success=True, message='该测试批次已清理')
                await db.commit()
                return response

            cls._assert_closed(batch)
            if command.batch_name != batch.batch_name:
                raise ServiceWarning('', '测试批次名称不匹配')
            cls._assert_token(batch, command.preview_token, datetime.now())

            selection = await TestBatchPurgeDao.load_selection(db, batch_id, for_update=True)
            if selection_hash(selection) != batch.preview_scope_hash:
                raise ServiceWarning('', '预览范围已变化')
            counts = cls._counts(selection)
            total_count = sum(counts.values())
            blocking_reason = cls._blocking_reason(
                total_count,
                TestBatchConfig.test_batch_online_purge_max_rows,
            )
            if blocking_reason is not None:
                raise ServiceWarning('', blocking_reason)

            await TestBatchPurgeDao.delete_user_warehouse_links(db, selection['biz_warehouse'])
            await TestBatchPurgeDao.delete_user_site_links(db, selection['biz_customer_site'])
            for table_name in derive_batch_delete_order(Base.metadata):
                deleted = await TestBatchPurgeDao.delete_table_rows(db, table_name, batch_id)
                if deleted != counts[table_name]:
                    raise ServiceException(message=f'{table_name} 清理数量发生变化')
            await TestBatchDao.delete_batch_users(db, batch_id)

            now = datetime.now()
            batch.status = TestBatchStatus.PURGED.value
            batch.purged_by = user_id
            batch.purged_at = now
            batch.purge_summary = json.dumps(counts, ensure_ascii=False, sort_keys=True)
            batch.update_by = user_name
            batch.update_time = now
            await TestBatchDao.clear_preview(db, batch)
            response = CrudResponseModel(is_success=True, message='测试批次清理成功')
            await db.commit()
            return response
        except Exception:
            await db.rollback()
            raise