from datetime import datetime from typing import Any from sqlalchemy import delete, func, select from sqlalchemy.ext.asyncio import AsyncSession from module_admin.entity.do.user_do import SysUser from module_business.domain.enums import TestBatchStatus from module_business.entity.do.master_do import ( BizCustomer, BizCustomerSite, BizPallet, BizProductionLine, BizStrain, BizWarehouse, ) from module_business.entity.do.operation_do import ( BizFillingRecord, BizInboundItem, BizInboundOrder, BizPalletEvent, BizPalletInventory, BizReturnItem, BizReturnOrder, BizShipmentItem, BizShipmentOrder, BizUsageRecord, ) from module_business.entity.do.test_batch_do import BizTestBatch, BizTestBatchUser _BATCH_OWNED_MODELS: tuple[type[Any], ...] = ( BizStrain, BizCustomer, BizCustomerSite, BizProductionLine, BizWarehouse, BizPallet, BizInboundOrder, BizInboundItem, BizFillingRecord, BizShipmentOrder, BizShipmentItem, BizUsageRecord, BizReturnOrder, BizReturnItem, BizPalletInventory, BizPalletEvent, ) class TestBatchDao: """测试批次访问数据库操作层。""" @classmethod async def get_batch(cls, db: AsyncSession, batch_id: int) -> BizTestBatch | None: result = await db.execute(select(BizTestBatch).where(BizTestBatch.batch_id == batch_id)) return result.scalar_one_or_none() @classmethod async def get_batch_for_update(cls, db: AsyncSession, batch_id: int) -> BizTestBatch | None: result = await db.execute( select(BizTestBatch) .where(BizTestBatch.batch_id == batch_id) .with_for_update() .execution_options(populate_existing=True) ) return result.scalar_one_or_none() @classmethod async def user_is_authorized(cls, db: AsyncSession, batch_id: int, user_id: int) -> bool: result = await db.execute( select(BizTestBatchUser.batch_id).where( BizTestBatchUser.batch_id == batch_id, BizTestBatchUser.user_id == user_id, ) ) return result.scalar_one_or_none() is not None @classmethod async def list_available_batches( cls, db: AsyncSession, user_id: int, is_admin: bool, ) -> list[BizTestBatch]: statement = select(BizTestBatch).where(BizTestBatch.status == TestBatchStatus.ACTIVE.value) if not is_admin: statement = statement.join(BizTestBatchUser).where(BizTestBatchUser.user_id == user_id) result = await db.execute(statement.order_by(BizTestBatch.batch_id.desc())) return list(result.scalars().all()) @classmethod async def list_batches(cls, db: AsyncSession) -> list[BizTestBatch]: result = await db.execute(select(BizTestBatch).order_by(BizTestBatch.batch_id.desc())) return list(result.scalars().all()) @classmethod async def add_batch(cls, db: AsyncSession, batch: BizTestBatch) -> None: db.add(batch) @classmethod async def list_batch_user_ids(cls, db: AsyncSession, batch_id: int) -> list[int]: result = await db.execute( select(BizTestBatchUser.user_id) .where(BizTestBatchUser.batch_id == batch_id) .order_by(BizTestBatchUser.user_id) ) return list(result.scalars().all()) @classmethod async def list_active_user_ids(cls, db: AsyncSession, user_ids: list[int]) -> list[int]: if not user_ids: return [] result = await db.execute( select(SysUser.user_id) .where( SysUser.user_id.in_(user_ids), SysUser.status == '0', SysUser.del_flag == '0', ) .order_by(SysUser.user_id) ) return list(result.scalars().all()) @classmethod async def replace_batch_users( cls, db: AsyncSession, batch_id: int, user_ids: list[int], create_by: str, create_time: datetime, ) -> None: await db.execute(delete(BizTestBatchUser).where(BizTestBatchUser.batch_id == batch_id)) if user_ids: db.add_all( [ BizTestBatchUser( batch_id=batch_id, user_id=user_id, create_by=create_by, create_time=create_time, ) for user_id in user_ids ] ) @classmethod async def count_owned_rows(cls, db: AsyncSession, batch_id: int) -> dict[str, int]: counts: dict[str, int] = {} for model in _BATCH_OWNED_MODELS: result = await db.execute(select(func.count()).select_from(model).where(model.test_batch_id == batch_id)) counts[model.__tablename__] = result.scalar() or 0 return counts @classmethod async def store_preview( cls, db: AsyncSession, batch: BizTestBatch, *, action: str, token_hash: str, scope_hash: str, expires_at: datetime, update_by: str, update_time: datetime, ) -> None: """在已锁定的批次对象上保存一次性预览状态。""" batch.preview_action = action batch.preview_token_hash = token_hash batch.preview_scope_hash = scope_hash batch.preview_expires_at = expires_at batch.update_by = update_by batch.update_time = update_time @classmethod async def clear_preview( cls, db: AsyncSession, batch: BizTestBatch, ) -> None: """消费已锁定批次上的预览令牌与范围摘要。""" batch.preview_action = None batch.preview_token_hash = None batch.preview_scope_hash = None batch.preview_expires_at = None