""" Copy Data Pipeline - Simple data copying without strategy execution """ from pathlib import Path from typing import List, Any, Optional from ..common.collection import DataCollection from ..common.types import SyncAction, SyncStatus from ..common.registry import DomainRegistry from ..datasource import JsonlDataSource from ..engine import StateMachineConfig, StateMachineRuntime, apply_state_to_node class CopyDataPipeline: """ 通用数据复制 Pipeline，支持任意数据源之间的复制。 流程： 1. 从 source_datasource 加载所有数据到 source_collection 2. 将每个节点复制到 target_collection，设置 action=CREATE 3. 使用 target_datasource 保存所有节点 用途： - 数据迁移 (JSONL → JSONL) - API 数据下载 (API → JSONL) - 测试数据准备 - 初始化同步 支持两种初始化方式： 1. 传入数据源实例 (推荐): source_datasource + target_datasource 2. 传入目录路径 (兼容): source_dir + target_dir (自动创建 JsonlDataSource) """ def __init__( self, source_datasource: Any = None, target_datasource: Any = None, source_dir: Optional[Path | str] = None, target_dir: Optional[Path | str] = None, node_types: List[str] | None = None, load_order: List[str] | None = None, save_order: List[str] | None = None, ): """ 初始化 CopyDataPipeline Args: source_datasource: 源数据源实例 (ApiDataSource, JsonlDataSource 等) target_datasource: 目标数据源实例 source_dir: 源数据目录 (兼容模式，会创建 JsonlDataSource) target_dir: 目标数据目录 (兼容模式，会创建 JsonlDataSource) node_types: 要复制的节点类型列表，None 表示复制所有注册的类型 load_order: 数据加载顺序，None 表示使用 node_types 顺序 save_order: 数据保存顺序，None 表示使用 node_types 顺序 """ self._supported_types: Optional[List[str]] = None # 数据源处理 - 优先使用实例，其次使用目录 if source_datasource is not None: self.source_datasource = source_datasource self._owns_source = False elif source_dir is not None: from ..domain.registry import JSONL_HANDLERS self.source_datasource = JsonlDataSource(dir_path=Path(source_dir), read_only=True) # 自动注册所有 JSONL handlers，并收集支持的类型 supported_types = [] for handler_class in JSONL_HANDLERS: handler = handler_class(self.source_datasource) self.source_datasource.register_handler(handler) supported_types.append(handler.node_type) self._owns_source = True self._supported_types = supported_types else: raise ValueError("Must provide either source_datasource or source_dir") if target_datasource is not None: self.target_datasource = target_datasource self._owns_target = False elif target_dir is not None: from ..domain.registry import JSONL_HANDLERS self.target_datasource = JsonlDataSource(dir_path=Path(target_dir), read_only=False) # 自动注册所有 JSONL handlers for handler_class in JSONL_HANDLERS: self.target_datasource.register_handler(handler_class(self.target_datasource)) self._owns_target = True else: raise ValueError("Must provide either target_datasource or target_dir") # 当使用 source_dir 创建时，如果 node_types 未指定，只使用支持的类型 if node_types is None and self._supported_types is not None: self.node_types = self._supported_types else: self.node_types = node_types or DomainRegistry.get_registered_types() self.load_order = load_order or self.node_types self.save_order = save_order or self.node_types cfg_path = Path(__file__).resolve().parents[1] / "config" / "node_state_machine.yaml" self._sm_runtime = StateMachineRuntime(StateMachineConfig.load(cfg_path)) # 统计信息 self.stats = { "loaded": 0, "copied": 0, "saved": 0, "failed": 0, } async def run(self): """ 执行完整的复制流程 Returns: dict: 统计信息 {"loaded": int, "copied": int, "saved": int, "failed": int} """ print(f"[CopyDataPipeline] Starting...") print(f" Source: {self.source_datasource.__class__.__name__}") print(f" Target: {self.target_datasource.__class__.__name__}") print(f" Node types: {self.node_types}") # Stage 1: 加载源数据 print("\n[Stage 1] Loading source data...") source_collection = DataCollection(collection_id="source") await self._load_source_data(source_collection) # Stage 2: 复制节点 print("\n[Stage 2] Copying nodes...") target_collection = DataCollection(collection_id="target") await self._copy_nodes(source_collection, target_collection) # Stage 3: 保存到目标 print("\n[Stage 3] Saving to target...") await self._save_target_data(target_collection) # 打印统计信息 self._print_summary() return self.stats async def _load_source_data(self, collection: DataCollection): """从源数据源加载所有节点""" # 如果是 JsonlDataSource 且没有注册 handlers，自动注册 if isinstance(self.source_datasource, JsonlDataSource) and not self.source_datasource._handlers: self._auto_register_jsonl_handlers(self.source_datasource) self.source_datasource.set_collection(collection) await self.source_datasource.load_all(order=self.load_order) # 统计加载的节点数 for node_type in self.node_types: count = len(collection.filter(node_type=node_type)) if count > 0: print(f" [LOAD] {node_type}: {count} nodes") self.stats["loaded"] += count async def _copy_nodes( self, source_collection: DataCollection, target_collection: DataCollection ): """ 复制所有节点到目标 collection，设置 action=CREATE, status=PENDING 注意： - 保持相同的 node_id（用于追踪） - data_id 设为空字符串（待创建） - 复制 data 和 depend_ids """ for node in source_collection.filter(): # filter() 不带参数返回所有节点 # 获取节点类 node_class = DomainRegistry.get_node_class(node.node_type) # 创建新节点 new_node = node_class( node_id=node.node_id, # 保持相同的 node_id data_id="", # CREATE 节点的 data_id 初始为空 data=node.data, # 复制数据 depend_ids=node.depend_ids.copy() if node.depend_ids else [], ) apply_state_to_node( self._sm_runtime, new_node, state_id="S06", reason="copy pipeline init", fields={"binding_status", "action", "status"}, ) apply_state_to_node( self._sm_runtime, new_node, state_id="S06", reason="copy pipeline init", fields={"data_id"}, data_id="", ) await target_collection.add(new_node) self.stats["copied"] += 1 print(f" Copied {self.stats['copied']} nodes") async def _save_target_data(self, collection: DataCollection): """保存所有节点到目标数据源""" # 如果是 JsonlDataSource 且没有注册 handlers，自动注册 if isinstance(self.target_datasource, JsonlDataSource) and not self.target_datasource._handlers: self._auto_register_jsonl_handlers(self.target_datasource) self.target_datasource.set_collection(collection) await self.target_datasource.sync_all(order=self.save_order) # 统计保存结果 for node_type in self.node_types: nodes = collection.filter(node_type=node_type) success_count = sum(1 for n in nodes if n.status == SyncStatus.SUCCESS) failed_count = sum(1 for n in nodes if n.status == SyncStatus.FAILED) if success_count > 0 or failed_count > 0: print(f" [SAVE] {node_type}: {success_count} success, {failed_count} failed") self.stats["saved"] += success_count self.stats["failed"] += failed_count def _auto_register_jsonl_handlers(self, datasource: JsonlDataSource): """自动为 JsonlDataSource 注册需要的 handlers（兼容模式）""" for node_type in self.node_types: if not DomainRegistry.is_registered(node_type): continue handler_class = DomainRegistry.get_jsonl_handler(node_type) if handler_class: handler = handler_class(datasource) datasource.register_handler(handler) def _print_summary(self): """打印执行摘要""" print("\n" + "=" * 60) print("[CopyDataPipeline] Execution Summary") print("=" * 60) print(f" Loaded: {self.stats['loaded']} nodes") print(f" Copied: {self.stats['copied']} nodes") print(f" Saved: {self.stats['saved']} nodes") print(f" Failed: {self.stats['failed']} nodes") print("=" * 60) if self.stats["failed"] > 0: print(f"\n⚠️ {self.stats['failed']} nodes failed to save") else: print("\n✅ All nodes saved successfully")