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

from sqlalchemy import or_, select
from sqlalchemy.exc import MultipleResultsFound
from sqlalchemy.ext.asyncio import AsyncSession

from exceptions.exception import ServiceException
from module_business.domain.test_batch_scope import TestBatchScope, batch_scope_condition
from module_business.entity.do.master_do import BizPallet
from module_business.entity.do.operation_do import BizPalletEvent, BizPalletInventory


class InventoryDao:
    """托盘库存锁、资产锁和只追加事件数据库操作层。"""

    @classmethod
    async def get_inventory_for_update(
        cls,
        db: AsyncSession,
        pallet_id: int,
        scope: TestBatchScope,
    ) -> BizPalletInventory | None:
        result = await db.execute(
            select(BizPalletInventory)
            .where(
                BizPalletInventory.pallet_id == pallet_id,
                batch_scope_condition(BizPalletInventory.test_batch_id, scope),
            )
            .with_for_update()
            .execution_options(populate_existing=True)
        )
        return result.scalar_one_or_none()

    @classmethod
    async def get_inventories_for_update(
        cls,
        db: AsyncSession,
        pallet_ids: list[int],
        scope: TestBatchScope,
    ) -> list[BizPalletInventory]:
        ordered_ids = sorted(set(pallet_ids))
        if not ordered_ids:
            return []

        result = await db.execute(
            select(BizPalletInventory)
            .where(
                BizPalletInventory.pallet_id.in_(ordered_ids),
                batch_scope_condition(BizPalletInventory.test_batch_id, scope),
            )
            .order_by(BizPalletInventory.pallet_id)
            .with_for_update()
            .execution_options(populate_existing=True)
        )
        return sorted(result.scalars().all(), key=lambda inventory: inventory.pallet_id)

    @classmethod
    async def get_pallet_by_qr_token(
        cls,
        db: AsyncSession,
        qr_token: str,
        scope: TestBatchScope,
    ) -> BizPallet | None:
        """兼容印刷托盘号；调用方仍须复核仓库/现场权限，写入使用规范二维码。"""
        identifier = qr_token.strip()
        result = await db.execute(
            select(BizPallet).where(
                or_(BizPallet.qr_token == identifier, BizPallet.pallet_code == identifier),
                batch_scope_condition(BizPallet.test_batch_id, scope),
            )
        )
        try:
            return result.scalar_one_or_none()
        except MultipleResultsFound:
            raise ServiceException(message='托盘号与二维码标识存在歧义，请联系管理员核对后重试') from None

    @classmethod
    async def get_pallets_by_ids(
        cls,
        db: AsyncSession,
        pallet_ids: list[int],
        scope: TestBatchScope,
    ) -> list[BizPallet]:
        """在调用方已锁库存后，按托盘 ID 升序锁定资产行。"""
        ordered_ids = sorted(set(pallet_ids))
        if not ordered_ids:
            return []

        result = await db.execute(
            select(BizPallet)
            .where(
                BizPallet.pallet_id.in_(ordered_ids),
                batch_scope_condition(BizPallet.test_batch_id, scope),
            )
            .order_by(BizPallet.pallet_id)
            .with_for_update()
            .execution_options(populate_existing=True)
        )
        return sorted(result.scalars().all(), key=lambda pallet: pallet.pallet_id)

    @classmethod
    async def get_event_by_idempotency_key(
        cls,
        db: AsyncSession,
        pallet_id: int,
        idempotency_key: str,
        scope: TestBatchScope,
    ) -> BizPalletEvent | None:
        result = await db.execute(
            select(BizPalletEvent).where(
                BizPalletEvent.pallet_id == pallet_id,
                BizPalletEvent.idempotency_key == idempotency_key,
                batch_scope_condition(BizPalletEvent.test_batch_id, scope),
            )
        )
        return result.scalar_one_or_none()

    @classmethod
    async def add_event(
        cls,
        db: AsyncSession,
        inventory: BizPalletInventory,
        event: BizPalletEvent,
        scope: TestBatchScope,
    ) -> BizPalletEvent:
        """在调用方持有库存行锁时，为托盘追加下一序号事件。"""
        last_sequence = 0
        if inventory.last_event_id is not None:
            sequence_result = await db.execute(
                select(BizPalletEvent.event_seq)
                .where(
                    BizPalletEvent.event_id == inventory.last_event_id,
                    BizPalletEvent.pallet_id == inventory.pallet_id,
                    batch_scope_condition(BizPalletEvent.test_batch_id, scope),
                )
                .with_for_update()
            )
            last_sequence = sequence_result.scalar_one_or_none()
            if last_sequence is None:
                raise RuntimeError('托盘库存事件指针不一致')
        event.pallet_id = inventory.pallet_id
        event.event_seq = last_sequence + 1
        db.add(event)
        await db.flush()
        inventory.last_event_id = event.event_id
        await db.flush()
        return event