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

This commit is contained in:
焦龙言
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)