first commit
This commit is contained in:
strepsiades
@@ -0,0 +1,247 @@
|
||||
"""
|
||||
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")
|
||||
Reference in New Issue
Block a user