import hashlib import json import secrets from collections import defaultdict, deque from dataclasses import dataclass from sqlalchemy import Column, select from sqlalchemy.ext.asyncio import AsyncSession from config.database import Base 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.vo.test_batch_vo import AdoptionRootModel BatchSelection = dict[str, frozenset[int]] BATCH_MODELS = { 'biz_strain': (BizStrain, BizStrain.strain_id, BizStrain.strain_code), 'biz_customer': (BizCustomer, BizCustomer.customer_id, BizCustomer.customer_code), 'biz_customer_site': (BizCustomerSite, BizCustomerSite.site_id, BizCustomerSite.site_code), 'biz_production_line': ( BizProductionLine, BizProductionLine.line_id, BizProductionLine.line_code, ), 'biz_warehouse': (BizWarehouse, BizWarehouse.warehouse_id, BizWarehouse.warehouse_code), 'biz_pallet': (BizPallet, BizPallet.pallet_id, BizPallet.pallet_code), 'biz_inbound_order': (BizInboundOrder, BizInboundOrder.inbound_id, BizInboundOrder.inbound_no), 'biz_inbound_item': ( BizInboundItem, BizInboundItem.inbound_item_id, BizInboundItem.inbound_item_id, ), 'biz_filling_record': ( BizFillingRecord, BizFillingRecord.filling_id, BizFillingRecord.filling_no, ), 'biz_shipment_order': ( BizShipmentOrder, BizShipmentOrder.shipment_id, BizShipmentOrder.shipment_no, ), 'biz_shipment_item': ( BizShipmentItem, BizShipmentItem.shipment_item_id, BizShipmentItem.shipment_item_id, ), 'biz_usage_record': (BizUsageRecord, BizUsageRecord.usage_id, BizUsageRecord.usage_no), 'biz_return_order': (BizReturnOrder, BizReturnOrder.return_id, BizReturnOrder.return_no), 'biz_return_item': ( BizReturnItem, BizReturnItem.return_item_id, BizReturnItem.return_item_id, ), 'biz_pallet_inventory': ( BizPalletInventory, BizPalletInventory.pallet_id, BizPalletInventory.pallet_id, ), 'biz_pallet_event': (BizPalletEvent, BizPalletEvent.event_id, BizPalletEvent.event_id), } BATCH_LABELS = { 'biz_strain': '菌种', 'biz_customer': '客户', 'biz_customer_site': '客户现场', 'biz_production_line': '生产线', 'biz_warehouse': '仓库', 'biz_pallet': '托盘', 'biz_inbound_order': '入库单', 'biz_inbound_item': '入库明细', 'biz_filling_record': '装填记录', 'biz_shipment_order': '发货单', 'biz_shipment_item': '发货明细', 'biz_usage_record': '使用记录', 'biz_return_order': '退回单', 'biz_return_item': '退回明细', 'biz_pallet_inventory': '托盘库存', 'biz_pallet_event': '托盘事件', } ROOT_TABLES = { 'strain': 'biz_strain', 'customer': 'biz_customer', 'site': 'biz_customer_site', 'line': 'biz_production_line', 'warehouse': 'biz_warehouse', 'pallet': 'biz_pallet', 'inbound': 'biz_inbound_order', 'filling': 'biz_filling_record', 'shipment': 'biz_shipment_order', 'usage': 'biz_usage_record', 'return': 'biz_return_order', } _EVENT_SOURCE_TABLES = { 'INBOUND': 'biz_inbound_order', 'FILLING': 'biz_filling_record', 'SHIPMENT': 'biz_shipment_order', 'USAGE': 'biz_usage_record', 'RETURN': 'biz_return_order', } def canonical_selection(selection: BatchSelection) -> bytes: payload = {table: sorted(ids) for table, ids in sorted(selection.items()) if ids} return json.dumps(payload, ensure_ascii=True, separators=(',', ':'), sort_keys=True).encode() def selection_hash(selection: BatchSelection) -> str: return hashlib.sha256(canonical_selection(selection)).hexdigest() def issue_preview_token() -> tuple[str, str]: raw = secrets.token_urlsafe(32) return raw, hashlib.sha256(raw.encode()).hexdigest() @dataclass(frozen=True) class _GraphEdge: source_table: str target_table: str filter_column: Column result_column: Column source_type: str | None = None class TestBatchGraph: """基于批次模型外键计算可原子纳入的双向关联闭包。""" @staticmethod def _foreign_key_edges() -> dict[str, list[_GraphEdge]]: edges: dict[str, list[_GraphEdge]] = defaultdict(list) batch_tables = set(BATCH_MODELS) for child_table_name in sorted(batch_tables): child_table = Base.metadata.tables[child_table_name] child_pk = BATCH_MODELS[child_table_name][1] for foreign_key in sorted( child_table.foreign_keys, key=lambda item: (item.parent.name, item.target_fullname), ): parent_table_name = foreign_key.column.table.name if foreign_key.parent.name == 'test_batch_id' or parent_table_name not in batch_tables: continue edges[child_table_name].append( _GraphEdge( source_table=child_table_name, target_table=parent_table_name, filter_column=child_pk, result_column=foreign_key.parent, ) ) edges[parent_table_name].append( _GraphEdge( source_table=parent_table_name, target_table=child_table_name, filter_column=foreign_key.parent, result_column=child_pk, ) ) return edges @classmethod def _all_edges(cls) -> dict[str, list[_GraphEdge]]: edges = cls._foreign_key_edges() event_pk = BATCH_MODELS['biz_pallet_event'][1] for source_type, root_table in sorted(_EVENT_SOURCE_TABLES.items()): edges[root_table].append( _GraphEdge( source_table=root_table, target_table='biz_pallet_event', filter_column=BizPalletEvent.source_id, result_column=event_pk, source_type=source_type, ) ) edges['biz_pallet_event'].append( _GraphEdge( source_table='biz_pallet_event', target_table=root_table, filter_column=event_pk, result_column=BizPalletEvent.source_id, source_type=source_type, ) ) for table_edges in edges.values(): table_edges.sort( key=lambda edge: ( edge.target_table, edge.filter_column.name, edge.result_column.name, edge.source_type or '', ) ) return edges @staticmethod async def _existing_ids( db: AsyncSession, table_name: str, ids: set[int], *, for_update: bool, ) -> set[int]: if not ids: return set() primary_key = BATCH_MODELS[table_name][1] statement = select(primary_key).where(primary_key.in_(sorted(ids))).order_by(primary_key) if for_update: statement = statement.with_for_update() result = await db.execute(statement) return {row[0] for row in result.all()} @classmethod async def _compute_closure( cls, db: AsyncSession, roots: list[AdoptionRootModel], *, for_update: bool, ) -> BatchSelection: selected: dict[str, set[int]] = {table_name: set() for table_name in BATCH_MODELS} requested: dict[str, set[int]] = defaultdict(set) for root in roots: requested[ROOT_TABLES[root.resource_type]].update(root.resource_ids) queue: deque[str] = deque() for table_name in sorted(requested): existing = await cls._existing_ids( db, table_name, requested[table_name], for_update=for_update, ) selected[table_name].update(existing) if existing: queue.append(table_name) expanded: dict[str, set[int]] = {table_name: set() for table_name in BATCH_MODELS} edges = cls._all_edges() while queue: source_table = queue.popleft() pending_ids = selected[source_table] - expanded[source_table] if not pending_ids: continue expanded[source_table].update(pending_ids) for edge in edges.get(source_table, []): statement = select(edge.result_column).where(edge.filter_column.in_(sorted(pending_ids))) if edge.source_type is not None: statement = statement.where(BizPalletEvent.source_type == edge.source_type) statement = statement.order_by(edge.result_column) if for_update: statement = statement.with_for_update() result = await db.execute(statement) related_ids = {row[0] for row in result.all() if row[0] is not None} new_ids = related_ids - selected[edge.target_table] if new_ids: selected[edge.target_table].update(new_ids) queue.append(edge.target_table) return {table_name: frozenset(selected[table_name]) for table_name in BATCH_MODELS} @classmethod async def compute_closure( cls, db: AsyncSession, roots: list[AdoptionRootModel], ) -> BatchSelection: return await cls._compute_closure(db, roots, for_update=False) @classmethod async def compute_locked_closure( cls, db: AsyncSession, roots: list[AdoptionRootModel], ) -> BatchSelection: """使用锁定读计算执行阶段闭包，避免复用一致性快照。""" return await cls._compute_closure(db, roots, for_update=True)