77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
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)
|