from datetime import datetime from typing import Any from pydantic import BaseModel from sqlalchemy import ColumnElement, Select, func, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm.attributes import InstrumentedAttribute from common.vo import PageModel from module_admin.entity.do.user_do import SysUser from module_business.domain.enums import PalletFlowState from module_business.domain.test_batch_scope import TestBatchScope, batch_scope_condition from module_business.entity.do.master_do import ( BizCustomer, BizCustomerSite, BizPallet, BizProductionLine, BizStrain, BizTunnel, BizUserSite, BizUserWarehouse, BizWarehouse, ) from module_business.entity.do.operation_do import BizPalletInventory from module_business.entity.vo.master_vo import ( CustomerModel, CustomerPageQueryModel, LineModel, LinePageQueryModel, PalletEditModel, PalletPageQueryModel, SiteModel, SitePageQueryModel, StrainModel, StrainPageQueryModel, TunnelModel, TunnelPageQueryModel, WarehouseModel, WarehousePageQueryModel, ) from utils.page_util import PageUtil _PALLET_FLOW_GROUP_STATES = { 'unstocked': (PalletFlowState.UNSTOCKED.value,), 'warehouse': (PalletFlowState.EMPTY_IN_WAREHOUSE.value, PalletFlowState.FILLED_IN_WAREHOUSE.value), 'transit': (PalletFlowState.IN_TRANSIT.value,), 'site': ( PalletFlowState.AT_SITE_UNUSED.value, PalletFlowState.AT_SITE_PARTIAL.value, PalletFlowState.AT_SITE_EMPTY.value, ), 'returning': (PalletFlowState.RETURNING.value,), } class MasterDao: """菌种、客户、现场、产线、隧道和仓库主数据数据库操作层。""" @staticmethod def _parse_id_list(ids_str: str | None) -> list[int]: if not ids_str: return [] result: list[int] = [] for item in ids_str.split(','): item = item.strip() if item.isdigit(): val = int(item) if val > 0: result.append(val) return result @staticmethod async def _paginate( db: AsyncSession, statement: Select, query_object: BaseModel, is_page: bool, ) -> PageModel | list[dict[str, Any]]: return await PageUtil.paginate( db, statement, query_object.page_num, query_object.page_size, is_page, exclude_fields={'test_batch_id'}, ) @staticmethod async def _get( db: AsyncSession, model: type[Any], id_column: InstrumentedAttribute, resource_id: int, scope_predicate: ColumnElement[bool], *, for_update: bool = False, ) -> Any | None: statement = select(model).where(id_column == resource_id, scope_predicate) if for_update: statement = statement.with_for_update().execution_options(populate_existing=True) result = await db.execute(statement) return result.scalar_one_or_none() @staticmethod async def _find( db: AsyncSession, model: type[Any], code_column: InstrumentedAttribute, code: str, *scope_predicates: ColumnElement[bool], ) -> Any | None: result = await db.execute(select(model).where(code_column == code, *scope_predicates)) return result.scalar_one_or_none() @staticmethod async def _add( db: AsyncSession, model: type[Any], id_field: str, page_object: BaseModel, test_batch_id: int | None, ) -> Any: payload = page_object.model_dump(exclude={id_field, 'can_remove', 'test_batch_id'}) payload['test_batch_id'] = test_batch_id record = model(**payload) db.add(record) await db.flush() return record @staticmethod async def _update( db: AsyncSession, model: type[Any], id_column: InstrumentedAttribute, id_field: str, page_object: BaseModel, scope: TestBatchScope, ) -> None: payload = page_object.model_dump( exclude_unset=True, exclude={'can_remove', 'create_by', 'create_time', 'test_batch_id'}, ) resource_id = payload.pop(id_field) await db.execute( update(model) .where( id_column == resource_id, batch_scope_condition(model.test_batch_id, scope), ) .values(**payload) ) @classmethod async def list_strain( cls, db: AsyncSession, query_object: StrainPageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizStrain.test_batch_id, scope)] if query_object.strain_code: conditions.append(BizStrain.strain_code.like(f'%{query_object.strain_code}%')) if query_object.strain_name: conditions.append(BizStrain.strain_name.like(f'%{query_object.strain_name}%')) if query_object.status: conditions.append(BizStrain.status == query_object.status) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizStrain.strain_id.in_(ids)) statement = select(BizStrain).where(*conditions).order_by(BizStrain.strain_id) return await cls._paginate(db, statement, query_object, is_page) @classmethod async def get_strain( cls, db: AsyncSession, strain_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizStrain | None: return await cls._get( db, BizStrain, BizStrain.strain_id, strain_id, batch_scope_condition(BizStrain.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_strain_by_code(cls, db: AsyncSession, strain_code: str) -> BizStrain | None: return await cls._find(db, BizStrain, BizStrain.strain_code, strain_code) @classmethod async def add_strain(cls, db: AsyncSession, strain: StrainModel, test_batch_id: int | None) -> BizStrain: return await cls._add(db, BizStrain, 'strain_id', strain, test_batch_id) @classmethod async def update_strain(cls, db: AsyncSession, strain: StrainModel, scope: TestBatchScope) -> None: await cls._update(db, BizStrain, BizStrain.strain_id, 'strain_id', strain, scope) @classmethod async def list_customer( cls, db: AsyncSession, query_object: CustomerPageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizCustomer.test_batch_id, scope)] if query_object.customer_code: conditions.append(BizCustomer.customer_code.like(f'%{query_object.customer_code}%')) if query_object.customer_name: conditions.append(BizCustomer.customer_name.like(f'%{query_object.customer_name}%')) if query_object.status: conditions.append(BizCustomer.status == query_object.status) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizCustomer.customer_id.in_(ids)) statement = select(BizCustomer).where(*conditions).order_by(BizCustomer.customer_id) return await cls._paginate(db, statement, query_object, is_page) @classmethod async def get_customer( cls, db: AsyncSession, customer_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizCustomer | None: return await cls._get( db, BizCustomer, BizCustomer.customer_id, customer_id, batch_scope_condition(BizCustomer.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_customer_by_code(cls, db: AsyncSession, customer_code: str) -> BizCustomer | None: return await cls._find(db, BizCustomer, BizCustomer.customer_code, customer_code) @classmethod async def add_customer(cls, db: AsyncSession, customer: CustomerModel, test_batch_id: int | None) -> BizCustomer: return await cls._add(db, BizCustomer, 'customer_id', customer, test_batch_id) @classmethod async def update_customer(cls, db: AsyncSession, customer: CustomerModel, scope: TestBatchScope) -> None: await cls._update(db, BizCustomer, BizCustomer.customer_id, 'customer_id', customer, scope) @classmethod async def list_site( cls, db: AsyncSession, query_object: SitePageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizCustomerSite.test_batch_id, scope)] if query_object.customer_id: conditions.append(BizCustomerSite.customer_id == query_object.customer_id) if query_object.site_code: conditions.append(BizCustomerSite.site_code.like(f'%{query_object.site_code}%')) if query_object.site_name: conditions.append(BizCustomerSite.site_name.like(f'%{query_object.site_name}%')) if query_object.status: conditions.append(BizCustomerSite.status == query_object.status) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizCustomerSite.site_id.in_(ids)) statement = select(BizCustomerSite).where(*conditions).order_by(BizCustomerSite.site_id) return await cls._paginate(db, statement, query_object, is_page) @classmethod async def get_site( cls, db: AsyncSession, site_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizCustomerSite | None: return await cls._get( db, BizCustomerSite, BizCustomerSite.site_id, site_id, batch_scope_condition(BizCustomerSite.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_site_by_code( cls, db: AsyncSession, site_code: str, customer_id: int, ) -> BizCustomerSite | None: return await cls._find( db, BizCustomerSite, BizCustomerSite.site_code, site_code, BizCustomerSite.customer_id == customer_id, ) @classmethod async def add_site(cls, db: AsyncSession, site: SiteModel, test_batch_id: int | None) -> BizCustomerSite: return await cls._add(db, BizCustomerSite, 'site_id', site, test_batch_id) @classmethod async def update_site(cls, db: AsyncSession, site: SiteModel, scope: TestBatchScope) -> None: await cls._update(db, BizCustomerSite, BizCustomerSite.site_id, 'site_id', site, scope) @classmethod async def list_line( cls, db: AsyncSession, query_object: LinePageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizProductionLine.test_batch_id, scope)] if query_object.site_id: conditions.append(BizProductionLine.site_id == query_object.site_id) if query_object.line_code: conditions.append(BizProductionLine.line_code.like(f'%{query_object.line_code}%')) if query_object.line_name: conditions.append(BizProductionLine.line_name.like(f'%{query_object.line_name}%')) if query_object.status: conditions.append(BizProductionLine.status == query_object.status) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizProductionLine.line_id.in_(ids)) statement = select(BizProductionLine).where(*conditions).order_by(BizProductionLine.line_id) return await cls._paginate(db, statement, query_object, is_page) @classmethod async def get_line( cls, db: AsyncSession, line_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizProductionLine | None: return await cls._get( db, BizProductionLine, BizProductionLine.line_id, line_id, batch_scope_condition(BizProductionLine.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_line_by_code( cls, db: AsyncSession, line_code: str, site_id: int, ) -> BizProductionLine | None: return await cls._find( db, BizProductionLine, BizProductionLine.line_code, line_code, BizProductionLine.site_id == site_id, ) @classmethod async def add_line(cls, db: AsyncSession, line: LineModel, test_batch_id: int | None) -> BizProductionLine: return await cls._add(db, BizProductionLine, 'line_id', line, test_batch_id) @classmethod async def update_line(cls, db: AsyncSession, line: LineModel, scope: TestBatchScope) -> None: await cls._update(db, BizProductionLine, BizProductionLine.line_id, 'line_id', line, scope) @classmethod async def list_tunnel( cls, db: AsyncSession, query_object: TunnelPageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizTunnel.test_batch_id, scope)] if query_object.site_id: conditions.append(BizTunnel.site_id == query_object.site_id) if query_object.tunnel_code: conditions.append(BizTunnel.tunnel_code.like(f'%{query_object.tunnel_code}%')) if query_object.tunnel_name: conditions.append(BizTunnel.tunnel_name.like(f'%{query_object.tunnel_name}%')) if query_object.status: conditions.append(BizTunnel.status == query_object.status) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizTunnel.tunnel_id.in_(ids)) statement = select(BizTunnel).where(*conditions).order_by(BizTunnel.tunnel_id) return await cls._paginate(db, statement, query_object, is_page) @classmethod async def get_tunnel( cls, db: AsyncSession, tunnel_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizTunnel | None: return await cls._get( db, BizTunnel, BizTunnel.tunnel_id, tunnel_id, batch_scope_condition(BizTunnel.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_tunnel_by_code( cls, db: AsyncSession, tunnel_code: str, site_id: int, ) -> BizTunnel | None: return await cls._find( db, BizTunnel, BizTunnel.tunnel_code, tunnel_code, BizTunnel.site_id == site_id, ) @classmethod async def add_tunnel(cls, db: AsyncSession, tunnel: TunnelModel, test_batch_id: int | None) -> BizTunnel: return await cls._add(db, BizTunnel, 'tunnel_id', tunnel, test_batch_id) @classmethod async def update_tunnel(cls, db: AsyncSession, tunnel: TunnelModel, scope: TestBatchScope) -> None: await cls._update(db, BizTunnel, BizTunnel.tunnel_id, 'tunnel_id', tunnel, scope) @classmethod async def list_warehouse( cls, db: AsyncSession, query_object: WarehousePageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizWarehouse.test_batch_id, scope)] if query_object.warehouse_code: conditions.append(BizWarehouse.warehouse_code.like(f'%{query_object.warehouse_code}%')) if query_object.warehouse_name: conditions.append(BizWarehouse.warehouse_name.like(f'%{query_object.warehouse_name}%')) if query_object.status: conditions.append(BizWarehouse.status == query_object.status) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizWarehouse.warehouse_id.in_(ids)) statement = select(BizWarehouse).where(*conditions).order_by(BizWarehouse.warehouse_id) return await cls._paginate(db, statement, query_object, is_page) @classmethod async def get_warehouse( cls, db: AsyncSession, warehouse_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizWarehouse | None: return await cls._get( db, BizWarehouse, BizWarehouse.warehouse_id, warehouse_id, batch_scope_condition(BizWarehouse.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_warehouse_by_code(cls, db: AsyncSession, warehouse_code: str) -> BizWarehouse | None: return await cls._find(db, BizWarehouse, BizWarehouse.warehouse_code, warehouse_code) @classmethod async def add_warehouse( cls, db: AsyncSession, warehouse: WarehouseModel, test_batch_id: int | None ) -> BizWarehouse: return await cls._add(db, BizWarehouse, 'warehouse_id', warehouse, test_batch_id) @classmethod async def update_warehouse(cls, db: AsyncSession, warehouse: WarehouseModel, scope: TestBatchScope) -> None: await cls._update(db, BizWarehouse, BizWarehouse.warehouse_id, 'warehouse_id', warehouse, scope) @classmethod async def list_pallet( cls, db: AsyncSession, query_object: PalletPageQueryModel, scope: TestBatchScope, is_page: bool = True, ) -> PageModel | list[dict[str, Any]]: conditions = [batch_scope_condition(BizPallet.test_batch_id, scope)] if query_object.pallet_code: # 扫码枪输入的是二维码固定标识，按编号模糊或标识精确匹配 conditions.append( or_( BizPallet.pallet_code.like(f'%{query_object.pallet_code}%'), BizPallet.qr_token == query_object.pallet_code, ) ) if query_object.asset_status: conditions.append(BizPallet.asset_status == query_object.asset_status) if query_object.flow_group: conditions.append(cls._pallet_flow_group_condition(query_object.flow_group, scope)) ids = cls._parse_id_list(getattr(query_object, 'ids', None)) if ids: conditions.append(BizPallet.pallet_id.in_(ids)) statement = select(BizPallet).where(*conditions).order_by(BizPallet.pallet_id) return await cls._paginate(db, statement, query_object, is_page) @staticmethod def _pallet_flow_group_condition(flow_group: str, scope: TestBatchScope) -> ColumnElement[bool]: states = _PALLET_FLOW_GROUP_STATES[flow_group] in_states = BizPallet.pallet_id.in_( select(BizPalletInventory.pallet_id).where( BizPalletInventory.flow_state.in_(states), batch_scope_condition(BizPalletInventory.test_batch_id, scope), ) ) if flow_group != 'unstocked': return in_states # 历史托盘可能缺少库存行，同样视为未入库 return or_( in_states, BizPallet.pallet_id.not_in( select(BizPalletInventory.pallet_id).where( batch_scope_condition(BizPalletInventory.test_batch_id, scope) ) ), ) @classmethod async def get_pallet( cls, db: AsyncSession, pallet_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizPallet | None: return await cls._get( db, BizPallet, BizPallet.pallet_id, pallet_id, batch_scope_condition(BizPallet.test_batch_id, scope), for_update=for_update, ) @classmethod async def find_pallet_by_code(cls, db: AsyncSession, pallet_code: str) -> BizPallet | None: return await cls._find(db, BizPallet, BizPallet.pallet_code, pallet_code) @classmethod async def add_pallet(cls, db: AsyncSession, pallet: BizPallet) -> BizPallet: db.add(pallet) await db.flush() return pallet @classmethod async def update_pallet( cls, db: AsyncSession, pallet: PalletEditModel, scope: TestBatchScope, update_by: str, update_time: datetime, ) -> None: await db.execute( update(BizPallet) .where( BizPallet.pallet_id == pallet.pallet_id, batch_scope_condition(BizPallet.test_batch_id, scope), ) .values( capacity_packages=pallet.capacity_packages, asset_status=pallet.asset_status, remark=pallet.remark, update_by=update_by, update_time=update_time, ) ) @classmethod async def get_pallet_inventory( cls, db: AsyncSession, pallet_id: int, scope: TestBatchScope, *, for_update: bool = False, ) -> BizPalletInventory | None: statement = select(BizPalletInventory).where( BizPalletInventory.pallet_id == pallet_id, batch_scope_condition(BizPalletInventory.test_batch_id, scope), ) if for_update: statement = statement.with_for_update().execution_options(populate_existing=True) result = await db.execute(statement) return result.scalar_one_or_none() @classmethod async def count_tunnel_trace_members( cls, db: AsyncSession, tunnel_id: int, trace_no: str, scope: TestBatchScope, ) -> int: result = await db.execute( select(func.count()) .select_from(BizPalletInventory) .where( BizPalletInventory.tunnel_id == tunnel_id, BizPalletInventory.trace_no == trace_no, batch_scope_condition(BizPalletInventory.test_batch_id, scope), ) ) return int(result.scalar_one() or 0) @classmethod async def add_pallet_inventory( cls, db: AsyncSession, inventory: BizPalletInventory, ) -> BizPalletInventory: db.add(inventory) await db.flush() return inventory @classmethod async def list_sites_by_code( cls, db: AsyncSession, site_code: str, scope: TestBatchScope, ) -> list[BizCustomerSite]: """按现场编码列出当前数据范围内的现场，用于导入时消除重名。""" return await cls._list_by_code(db, BizCustomerSite, BizCustomerSite.site_code, site_code, scope) @classmethod async def list_lines_by_code( cls, db: AsyncSession, line_code: str, scope: TestBatchScope, ) -> list[BizProductionLine]: """按产线编码列出当前数据范围内的产线，用于判断能否变更所属现场。""" return await cls._list_by_code(db, BizProductionLine, BizProductionLine.line_code, line_code, scope) @classmethod async def list_tunnels_by_code( cls, db: AsyncSession, tunnel_code: str, scope: TestBatchScope, ) -> list[BizTunnel]: """按隧道编号列出当前数据范围内的隧道，用于判断能否变更所属现场。""" return await cls._list_by_code(db, BizTunnel, BizTunnel.tunnel_code, tunnel_code, scope) @classmethod async def map_customer_codes(cls, db: AsyncSession, customer_ids: set[int]) -> dict[int, str]: """批量把客户 ID 转成客户编码，供导出模板列使用。""" if not customer_ids: return {} result = await db.execute( select(BizCustomer.customer_id, BizCustomer.customer_code).where(BizCustomer.customer_id.in_(customer_ids)) ) return {int(customer_id): customer_code for customer_id, customer_code in result.all()} @classmethod async def map_site_codes(cls, db: AsyncSession, site_ids: set[int]) -> dict[int, str]: """批量把现场 ID 转成现场编码，供产线和隧道导出使用。""" if not site_ids: return {} result = await db.execute( select(BizCustomerSite.site_id, BizCustomerSite.site_code).where(BizCustomerSite.site_id.in_(site_ids)) ) return {int(site_id): site_code for site_id, site_code in result.all()} @classmethod async def _list_by_code( cls, db: AsyncSession, model: type[Any], code_column: InstrumentedAttribute, code: str, scope: TestBatchScope, ) -> list[Any]: result = await db.execute( select(model).where(code_column == code, batch_scope_condition(model.test_batch_id, scope)) ) return list(result.scalars().all()) @classmethod async def list_assignable_users(cls, db: AsyncSession) -> list[Any]: result = await db.execute( select(SysUser.user_id, SysUser.user_name, SysUser.nick_name) .where(SysUser.status == '0', SysUser.del_flag == '0') .order_by(SysUser.user_id) ) return list(result.all()) @classmethod async def list_assignable_users_by_ids(cls, db: AsyncSession, user_ids: list[int]) -> list[Any]: if not user_ids: return [] result = await db.execute( select(SysUser.user_id, SysUser.user_name, SysUser.nick_name) .where( SysUser.user_id.in_(user_ids), SysUser.status == '0', SysUser.del_flag == '0', ) .order_by(SysUser.user_id) ) return list(result.all()) @classmethod async def list_warehouse_assigned_user_ids( cls, db: AsyncSession, warehouse_id: int, scope: TestBatchScope, ) -> list[int]: result = await db.execute( select(BizUserWarehouse.user_id) .join(BizWarehouse, BizWarehouse.warehouse_id == BizUserWarehouse.warehouse_id) .where( BizUserWarehouse.warehouse_id == warehouse_id, batch_scope_condition(BizWarehouse.test_batch_id, scope), ) ) return list(result.scalars().all()) @classmethod async def list_site_assigned_user_ids( cls, db: AsyncSession, site_id: int, scope: TestBatchScope, ) -> list[int]: result = await db.execute( select(BizUserSite.user_id) .join(BizCustomerSite, BizCustomerSite.site_id == BizUserSite.site_id) .where( BizUserSite.site_id == site_id, batch_scope_condition(BizCustomerSite.test_batch_id, scope), ) ) return list(result.scalars().all())