ShiErFeng/shierfeng-fastapi-backend/module_business/dao/test_batch_dao.py

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