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