import hashlib import heapq import json import secrets from collections import defaultdict, deque from dataclasses import dataclass from sqlalchemy import Column, MetaData, 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, BizTunnel, 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_tunnel': (BizTunnel, BizTunnel.tunnel_id, BizTunnel.tunnel_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_tunnel': '隧道', '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': '托盘事件', } _DELETE_TIE_BREAK_ORDER = ( 'biz_pallet_inventory', 'biz_usage_record', 'biz_return_item', 'biz_shipment_item', 'biz_inbound_item', 'biz_return_order', 'biz_shipment_order', 'biz_inbound_order', 'biz_pallet_event', 'biz_filling_record', 'biz_production_line', 'biz_tunnel', 'biz_customer_site', 'biz_customer', 'biz_warehouse', 'biz_strain', 'biz_pallet', ) ROOT_TABLES = { 'strain': 'biz_strain', 'customer': 'biz_customer', 'site': 'biz_customer_site', 'line': 'biz_production_line', 'tunnel': 'biz_tunnel', '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', 'TUNNEL': 'biz_tunnel', '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() def derive_batch_delete_order(metadata: MetaData) -> list[str]: """从实际外键图生成子表优先的确定性删除顺序。""" safe_tables = set(BATCH_MODELS) metadata_batch_tables = {table.name for table in metadata.tables.values() if 'test_batch_id' in table.columns} if metadata_batch_tables != safe_tables or set(_DELETE_TIE_BREAK_ORDER) != safe_tables: raise RuntimeError('测试批次清理安全表集合不一致') children_by_parent: dict[str, set[str]] = {table_name: set() for table_name in safe_tables} parents_by_child: dict[str, set[str]] = {table_name: set() for table_name in safe_tables} for child_name in sorted(safe_tables): for foreign_key in metadata.tables[child_name].foreign_keys: parent_name = foreign_key.target_fullname.rsplit('.', maxsplit=1)[0] if parent_name not in safe_tables: continue children_by_parent[parent_name].add(child_name) parents_by_child[child_name].add(parent_name) priority = {table_name: index for index, table_name in enumerate(_DELETE_TIE_BREAK_ORDER)} ready = [(priority[table_name], table_name) for table_name, children in children_by_parent.items() if not children] heapq.heapify(ready) order: list[str] = [] while ready: _, child_name = heapq.heappop(ready) order.append(child_name) for parent_name in sorted(parents_by_child[child_name]): children_by_parent[parent_name].remove(child_name) if not children_by_parent[parent_name]: heapq.heappush(ready, (priority[parent_name], parent_name)) if len(order) != len(safe_tables): raise RuntimeError('测试批次清理外键图存在环') if set(order) != safe_tables: raise RuntimeError('测试批次清理安全表集合不一致') return order @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)