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

shierfeng-fastapi-backend/module_business/dao/test_batch_purge_dao.py

@@ -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)