strepsiades
248 lines
10 KiB
Python
248 lines
10 KiB
Python
"""
|
|
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.add_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.add_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.add_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")
|