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

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

from sqlalchemy.ext.asyncio import AsyncSession

from common.vo import CrudResponseModel
from exceptions.exception import PermissionException, ServiceWarning
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_business.dao.test_batch_adoption_dao import TestBatchAdoptionDao
from module_business.dao.test_batch_dao import TestBatchDao
from module_business.domain.enums import TestBatchStatus
from module_business.domain.test_batch_graph import (
    BATCH_LABELS,
    BATCH_MODELS,
    TestBatchGraph,
    issue_preview_token,
    selection_hash,
)
from module_business.entity.vo.test_batch_vo import (
    AdoptionExecuteRequestModel,
    AdoptionPreviewRequestModel,
    BatchImpactModel,
    TestBatchPreviewModel,
)
from module_business.service.test_batch_access_service import TestBatchAccessService


class TestBatchAdoptionService:
    """预览并原子纳入已存在的正式数据闭包。"""

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

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

    @staticmethod
    def _assert_active(batch: Any | None) -> None:
        if batch is None:
            raise ServiceWarning('', '测试批次不存在')
        status = batch.status
        if status == TestBatchStatus.CLOSED.value:
            raise ServiceWarning('', '测试批次已关闭')
        if status == TestBatchStatus.PURGED.value:
            raise ServiceWarning('', '测试批次已清理')
        if status != TestBatchStatus.ACTIVE.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) != 'ADOPT':
            raise ServiceWarning('', '预览令牌无效')
        expires_at = getattr(batch, 'preview_expires_at', None)
        if expires_at is None or expires_at <= now:
            raise ServiceWarning('', '预览令牌已过期')

    @classmethod
    async def preview(
        cls,
        db: AsyncSession,
        batch_id: int,
        command: AdoptionPreviewRequestModel,
        current_user: CurrentUserModel,
    ) -> TestBatchPreviewModel:
        cls._assert_admin(current_user)
        user_name = cls._user_name(current_user)
        try:
            batch = await TestBatchDao.get_batch_for_update(db, batch_id)
            cls._assert_active(batch)
            selection = await TestBatchGraph.compute_closure(db, command.roots)
            if await TestBatchAdoptionDao.find_tagged_record(db, selection) is not None:
                raise ServiceWarning('', '存在记录已属于其他测试批次')
            display_keys = await TestBatchAdoptionDao.list_display_keys(db, selection)

            now = datetime.now()
            expires_at = now + timedelta(seconds=600)
            raw_token, token_hash = issue_preview_token()
            await TestBatchDao.store_preview(
                db,
                batch,
                action='ADOPT',
                token_hash=token_hash,
                scope_hash=selection_hash(selection),
                expires_at=expires_at,
                update_by=user_name,
                update_time=now,
            )
            await db.commit()
            impacts = [
                BatchImpactModel(
                    table_name=table_name,
                    label=BATCH_LABELS[table_name],
                    count=len(selection[table_name]),
                    keys=display_keys[table_name],
                )
                for table_name in BATCH_MODELS
            ]
            return TestBatchPreviewModel(
                batch_id=batch_id,
                batch_name=batch.batch_name,
                preview_token=raw_token,
                expires_at=expires_at,
                total_count=sum(item.count for item in impacts),
                impacts=impacts,
            )
        except Exception:
            await db.rollback()
            raise

    @classmethod
    async def execute(
        cls,
        db: AsyncSession,
        batch_id: int,
        command: AdoptionExecuteRequestModel,
        current_user: CurrentUserModel,
    ) -> CrudResponseModel:
        cls._assert_admin(current_user)
        user_name = cls._user_name(current_user)
        try:
            batch = await TestBatchDao.get_batch_for_update(db, batch_id)
            cls._assert_active(batch)
            if command.batch_name != batch.batch_name:
                raise ServiceWarning('', '测试批次名称不匹配')
            cls._assert_token(batch, command.preview_token, datetime.now())

            await TestBatchAdoptionDao.lock_roots(db, command.roots)
            first_selection = await TestBatchGraph.compute_closure(db, command.roots)
            await TestBatchAdoptionDao.lock_selection(db, first_selection)
            locked_selection = await TestBatchGraph.compute_closure(db, command.roots)
            expected_scope_hash = batch.preview_scope_hash
            if (
                selection_hash(first_selection) != expected_scope_hash
                or selection_hash(locked_selection) != expected_scope_hash
            ):
                raise ServiceWarning('', '预览范围已变化')
            if await TestBatchAdoptionDao.find_tagged_record(db, locked_selection) is not None:
                raise ServiceWarning('', '存在记录已属于其他测试批次')

            for table_name in sorted(locked_selection):
                await TestBatchAdoptionDao.set_test_batch_id(
                    db,
                    table_name,
                    sorted(locked_selection[table_name]),
                    batch_id,
                )
            now = datetime.now()
            await TestBatchDao.clear_preview(db, batch)
            batch.update_by = user_name
            batch.update_time = now
            await db.commit()
            return CrudResponseModel(is_success=True, message='旧测试数据纳入成功')
        except Exception:
            await db.rollback()
            raise