安全：增加测试批次原子清理
This commit is contained in:
焦龙言
@@ -0,0 +1,196 @@
|
||||
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
|
||||
Reference in New Issue
Block a user