安全：增加测试批次原子清理

2026-09-23 14:52:07 +08:00
parent b11c457e61
commit 1aaed57825
6 changed files with 966 additions and 1 deletions
@@ -145,6 +145,10 @@ class TestBatchDao:
                ]
            )

+    @classmethod
+    async def delete_batch_users(cls, db: AsyncSession, batch_id: int) -> None:
+        await db.execute(delete(BizTestBatchUser).where(BizTestBatchUser.batch_id == batch_id))
+
    @classmethod
    async def count_owned_rows(cls, db: AsyncSession, batch_id: int) -> dict[str, int]:
        counts: dict[str, int] = {}
@@ -0,0 +1,76 @@
+from collections.abc import Collection
+
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from exceptions.exception import ServiceException
+from module_business.domain.test_batch_graph import BATCH_MODELS, BatchSelection
+from module_business.entity.do.master_do import BizUserSite, BizUserWarehouse
+
+
+class TestBatchPurgeDao:
+    """测试批次清理范围查询与无提交删除。"""
+
+    @classmethod
+    async def load_selection(
+        cls,
+        db: AsyncSession,
+        batch_id: int,
+        *,
+        for_update: bool,
+    ) -> BatchSelection:
+        selection: BatchSelection = {}
+        for table_name, (model, primary_key, _) in BATCH_MODELS.items():
+            statement = select(primary_key).where(model.test_batch_id == batch_id).order_by(primary_key)
+            if for_update:
+                statement = statement.with_for_update()
+            result = await db.execute(statement)
+            selection[table_name] = frozenset(row[0] for row in result.all())
+        return selection
+
+    @classmethod
+    async def list_display_keys(
+        cls,
+        db: AsyncSession,
+        selection: BatchSelection,
+    ) -> dict[str, list[str]]:
+        keys: dict[str, list[str]] = {}
+        for table_name, (_, primary_key, display_column) in BATCH_MODELS.items():
+            ids = sorted(selection[table_name])
+            if not ids:
+                keys[table_name] = []
+                continue
+            result = await db.execute(select(display_column).where(primary_key.in_(ids)).order_by(display_column))
+            keys[table_name] = sorted(str(row[0]) for row in result.all())
+        return keys
+
+    @classmethod
+    async def delete_user_warehouse_links(
+        cls,
+        db: AsyncSession,
+        warehouse_ids: Collection[int],
+    ) -> None:
+        if warehouse_ids:
+            await db.execute(delete(BizUserWarehouse).where(BizUserWarehouse.warehouse_id.in_(sorted(warehouse_ids))))
+
+    @classmethod
+    async def delete_user_site_links(
+        cls,
+        db: AsyncSession,
+        site_ids: Collection[int],
+    ) -> None:
+        if site_ids:
+            await db.execute(delete(BizUserSite).where(BizUserSite.site_id.in_(sorted(site_ids))))
+
+    @classmethod
+    async def delete_table_rows(
+        cls,
+        db: AsyncSession,
+        table_name: str,
+        batch_id: int,
+    ) -> int:
+        if table_name not in BATCH_MODELS:
+            raise ServiceException(message=f'不允许清理非安全表: {table_name}')
+        model = BATCH_MODELS[table_name][0]
+        result = await db.execute(delete(model).where(model.test_batch_id == batch_id))
+        return int(result.rowcount or 0)
@@ -1,10 +1,11 @@
 import hashlib
+import heapq
 import json
 import secrets
 from collections import defaultdict, deque
 from dataclasses import dataclass

-from sqlalchemy import Column, select
+from sqlalchemy import Column, MetaData, select
 from sqlalchemy.ext.asyncio import AsyncSession

 from config.database import Base
@@ -98,6 +99,25 @@ BATCH_LABELS = {
    'biz_pallet_event': '托盘事件',
 }

+_DELETE_TIE_BREAK_ORDER = (
+    'biz_pallet_inventory',
+    'biz_usage_record',
+    'biz_return_item',
+    'biz_shipment_item',
+    'biz_inbound_item',
+    'biz_return_order',
+    'biz_shipment_order',
+    'biz_inbound_order',
+    'biz_pallet_event',
+    'biz_filling_record',
+    'biz_production_line',
+    'biz_customer_site',
+    'biz_customer',
+    'biz_warehouse',
+    'biz_strain',
+    'biz_pallet',
+)
+
 ROOT_TABLES = {
    'strain': 'biz_strain',
    'customer': 'biz_customer',
@@ -135,6 +155,42 @@ def issue_preview_token() -> tuple[str, str]:
    return raw, hashlib.sha256(raw.encode()).hexdigest()


+def derive_batch_delete_order(metadata: MetaData) -> list[str]:
+    """从实际外键图生成子表优先的确定性删除顺序。"""
+    safe_tables = set(BATCH_MODELS)
+    metadata_batch_tables = {table.name for table in metadata.tables.values() if 'test_batch_id' in table.columns}
+    if metadata_batch_tables != safe_tables or set(_DELETE_TIE_BREAK_ORDER) != safe_tables:
+        raise RuntimeError('测试批次清理安全表集合不一致')
+
+    children_by_parent: dict[str, set[str]] = {table_name: set() for table_name in safe_tables}
+    parents_by_child: dict[str, set[str]] = {table_name: set() for table_name in safe_tables}
+    for child_name in sorted(safe_tables):
+        for foreign_key in metadata.tables[child_name].foreign_keys:
+            parent_name = foreign_key.target_fullname.rsplit('.', maxsplit=1)[0]
+            if parent_name not in safe_tables:
+                continue
+            children_by_parent[parent_name].add(child_name)
+            parents_by_child[child_name].add(parent_name)
+
+    priority = {table_name: index for index, table_name in enumerate(_DELETE_TIE_BREAK_ORDER)}
+    ready = [(priority[table_name], table_name) for table_name, children in children_by_parent.items() if not children]
+    heapq.heapify(ready)
+    order: list[str] = []
+    while ready:
+        _, child_name = heapq.heappop(ready)
+        order.append(child_name)
+        for parent_name in sorted(parents_by_child[child_name]):
+            children_by_parent[parent_name].remove(child_name)
+            if not children_by_parent[parent_name]:
+                heapq.heappush(ready, (priority[parent_name], parent_name))
+
+    if len(order) != len(safe_tables):
+        raise RuntimeError('测试批次清理外键图存在环')
+    if set(order) != safe_tables:
+        raise RuntimeError('测试批次清理安全表集合不一致')
+    return order
+
+
@dataclass(frozen=True)
 class _GraphEdge:
    source_table: str
@@ -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