from datetime import datetime from typing import Any from pydantic import BaseModel from sqlalchemy import ColumnElement, Select, func, 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.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 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(BizPallet.pallet_code.like(f'%{query_object.pallet_code}%')) if query_object.asset_status: conditions.append(BizPallet.asset_status == query_object.asset_status) 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) @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())