first commit
This commit is contained in:
strepsiades
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
JSONL DataSource 模块
|
||||
|
||||
提供基于 JSONL 文件的数据源实现，用于测试和本地开发。
|
||||
"""
|
||||
|
||||
from .datasource import JsonlDataSource
|
||||
from .handler import BaseJsonlHandler
|
||||
|
||||
__all__ = [
|
||||
"JsonlDataSource",
|
||||
"BaseJsonlHandler",
|
||||
]
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
JsonlDataSource - JSONL 文件数据源
|
||||
|
||||
职责：
|
||||
- 提供 JSONL 文件目录管理
|
||||
- 实现 DataSource 接口（load_all, sync_all）
|
||||
- 管理内存缓存和文件写入
|
||||
- 统一异常处理和状态管理
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, Set
|
||||
|
||||
from ..datasource import BaseDataSource
|
||||
|
||||
|
||||
|
||||
|
||||
class JsonlDataSource(BaseDataSource):
|
||||
"""
|
||||
JSONL 文件数据源
|
||||
|
||||
提供文件路径和基础功能，具体的文件查找和读取由各个 Handler 负责。
|
||||
|
||||
Args:
|
||||
dir_path: JSONL 文件目录
|
||||
read_only: 是否只读模式
|
||||
|
||||
示例：
|
||||
datasource = JsonlDataSource(Path("filtered_datasource/datasource/ecm-2025-12-02"))
|
||||
datasource.register_handler(ContractJsonlHandler(datasource))
|
||||
|
||||
collection = DataCollection()
|
||||
datasource.set_collection(collection)
|
||||
await datasource.load_all(order=["project", "contract"])
|
||||
await datasource.sync_all(order=["project", "contract"])
|
||||
"""
|
||||
|
||||
def __init__(self, dir_path: Path, read_only: bool = False):
|
||||
super().__init__()
|
||||
self.dir_path = Path(dir_path)
|
||||
self.read_only = read_only
|
||||
|
||||
# 用于 create/update/delete 的内存缓存
|
||||
self._data: Dict[str, Dict[str, Dict[str, Any]]] = {} # {node_type: {id: data}}
|
||||
self._dirty_types: Set[str] = set()
|
||||
|
||||
def _on_node_loaded(self, handler, node) -> None:
|
||||
node.append_log(
|
||||
f"JSONL链路[load] dir={self.dir_path.as_posix()} node_type={node.node_type} data_id={node.data_id or '<empty>'}"
|
||||
)
|
||||
|
||||
def _on_in_progress_result(self, handler, node, result) -> None:
|
||||
node.append_log(
|
||||
f"JSONL链路[submit] action={node.action.value} data_id={node.data_id or '<empty>'}"
|
||||
)
|
||||
|
||||
def _on_poll_result(self, handler, node, task_id: str, result) -> None:
|
||||
node.append_log(
|
||||
f"JSONL链路[poll] task_id={task_id} status={result.status.value}"
|
||||
)
|
||||
|
||||
# load_all 和 sync_all 继承自 BaseDataSource
|
||||
# 保存逻辑在 BaseJsonlHandler 中实现
|
||||
|
||||
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
BaseJsonlHandler - JSONL Handler 基类
|
||||
|
||||
职责：
|
||||
- 查找和读取 JSONL 文件
|
||||
- 实现批量 Handler 接口（load, create_all, update_all, delete_all）
|
||||
- 使用 datasource 的内存缓存
|
||||
- 返回 List[TaskResult]（统一状态管理）
|
||||
- 捕获所有异常并转换为 TaskResult.failed()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Dict, List, Any, TYPE_CHECKING, Optional, Set
|
||||
from uuid import uuid4
|
||||
|
||||
from ..handler import BaseNodeHandler
|
||||
from ..task_result import TaskResult
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...common.sync_node import SyncNode
|
||||
from .datasource import JsonlDataSource
|
||||
|
||||
|
||||
class BaseJsonlHandler(BaseNodeHandler):
|
||||
"""
|
||||
JSONL Handler 基类
|
||||
|
||||
负责查找和读取对应的 JSONL 文件。
|
||||
|
||||
示例：
|
||||
class ContractJsonlHandler(BaseJsonlHandler):
|
||||
def __init__(self, datasource: JsonlDataSource):
|
||||
super().__init__(
|
||||
datasource=datasource,
|
||||
node_type="contract",
|
||||
node_class=ContractSyncNode,
|
||||
schema=ContractResponse
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
datasource: JsonlDataSource,
|
||||
node_type: str,
|
||||
node_class: type,
|
||||
schema: type
|
||||
):
|
||||
super().__init__(node_type, node_class, schema)
|
||||
self.datasource = datasource
|
||||
self._target_project_ids: Optional[Set[str]] = None
|
||||
|
||||
def set_target_project_ids(self, target_project_ids: List[str]) -> None:
|
||||
"""注入项目过滤列表，供 load 阶段预过滤使用。"""
|
||||
if target_project_ids:
|
||||
self._target_project_ids = {str(pid) for pid in target_project_ids if str(pid)}
|
||||
else:
|
||||
self._target_project_ids = None
|
||||
|
||||
def _should_skip_item_by_project_filter(self, item: Dict[str, Any], item_id: str) -> bool:
|
||||
if not self._target_project_ids:
|
||||
return False
|
||||
if self.node_type == "project":
|
||||
return item_id not in self._target_project_ids
|
||||
|
||||
project_id = item.get("project_id")
|
||||
if project_id:
|
||||
return str(project_id) not in self._target_project_ids
|
||||
return False
|
||||
|
||||
async def load(self) -> List['SyncNode']:
|
||||
"""
|
||||
查找并加载对应的 JSONL 文件并创建节点
|
||||
|
||||
文件命名规则：{node_type}_N.jsonl (N 为数字)
|
||||
例如：contract_1.jsonl, contract_2.jsonl
|
||||
|
||||
注意：load 操作的异常由 DataSource 的 load_all 捕获处理
|
||||
"""
|
||||
nodes = []
|
||||
|
||||
if not self.datasource.dir_path.exists():
|
||||
return nodes
|
||||
|
||||
# 查找所有匹配的文件：{node_type}_N.jsonl
|
||||
pattern = re.compile(rf"^{re.escape(self.node_type)}_\d+\.jsonl$")
|
||||
|
||||
for file_path in self.datasource.dir_path.iterdir():
|
||||
if not pattern.match(file_path.name):
|
||||
continue
|
||||
|
||||
# 读取文件
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(f, 1):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
item = json.loads(line)
|
||||
item_id = self.extract_id(item)
|
||||
|
||||
if self._should_skip_item_by_project_filter(item, item_id):
|
||||
continue
|
||||
|
||||
# 创建节点
|
||||
node = self.create_node(item)
|
||||
nodes.append(node)
|
||||
|
||||
# 同时加载到缓存
|
||||
if item_id:
|
||||
bucket = self.datasource._data.setdefault(self.node_type, {})
|
||||
bucket[item_id] = item
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Warning: Invalid JSON in {file_path.name}:{line_num}: {e}")
|
||||
# 继续处理下一行
|
||||
|
||||
return nodes
|
||||
|
||||
def extract_id(self, data: Dict[str, Any]) -> str:
|
||||
"""
|
||||
提取数据 ID
|
||||
|
||||
默认实现：优先使用 'id'，其次 '_id'
|
||||
子类可重写此方法。
|
||||
|
||||
Args:
|
||||
data: 数据字典
|
||||
|
||||
Returns:
|
||||
数据 ID
|
||||
"""
|
||||
return str(data.get("id") or data.get("_id") or "")
|
||||
|
||||
async def create_all(
|
||||
self,
|
||||
nodes: List['SyncNode']
|
||||
) -> List[TaskResult]:
|
||||
"""
|
||||
批量创建节点
|
||||
|
||||
Args:
|
||||
nodes: 需要创建的节点列表（DataSource 已筛选为 CREATE + PENDING）
|
||||
|
||||
Returns:
|
||||
List[TaskResult]，每个结果关联 node_id
|
||||
"""
|
||||
results = []
|
||||
|
||||
for node in nodes:
|
||||
try:
|
||||
# 生成新 ID（如果没有）
|
||||
if not node.data_id:
|
||||
new_id = str(uuid4())
|
||||
else:
|
||||
new_id = node.data_id
|
||||
|
||||
# 获取数据字典
|
||||
# 使用 exclude_unset=False 以确保测试中的完整性校验（如 amount 字段）能通过
|
||||
# 因为 JSONL 存储通常需要完整记录
|
||||
data_dict = node.get_data(exclude_unset=False)
|
||||
if data_dict is None:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error="Cannot get data from node"
|
||||
))
|
||||
continue
|
||||
|
||||
# 设置 ID 和 type
|
||||
data_dict["id"] = new_id
|
||||
data_dict["_id"] = new_id
|
||||
if "type" not in data_dict:
|
||||
data_dict["type"] = self.node_type
|
||||
|
||||
# 存储到缓存
|
||||
bucket = self.datasource._data.setdefault(self.node_type, {})
|
||||
bucket[new_id] = data_dict
|
||||
self.datasource._dirty_types.add(self.node_type)
|
||||
|
||||
results.append(TaskResult.success(
|
||||
node_id=node.node_id,
|
||||
data_id=new_id
|
||||
))
|
||||
|
||||
except Exception as e:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error=f"Create failed: {str(e)}"
|
||||
))
|
||||
|
||||
# 自动保存到文件
|
||||
await self._save_to_file()
|
||||
|
||||
return results
|
||||
|
||||
async def update_all(
|
||||
self,
|
||||
nodes: List['SyncNode']
|
||||
) -> List[TaskResult]:
|
||||
"""
|
||||
批量更新节点
|
||||
|
||||
Args:
|
||||
nodes: 需要更新的节点列表（DataSource 已筛选为 UPDATE + PENDING）
|
||||
|
||||
Returns:
|
||||
List[TaskResult]，每个结果关联 node_id
|
||||
"""
|
||||
results = []
|
||||
|
||||
for node in nodes:
|
||||
try:
|
||||
if not node.data_id:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error="Cannot update node without data_id"
|
||||
))
|
||||
continue
|
||||
|
||||
# 获取数据字典
|
||||
# JSONL 更新通常也是完整记录替换，所以使用 exclude_unset=False
|
||||
data_dict = node.get_data(exclude_unset=False)
|
||||
if data_dict is None:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error="Cannot get data from node"
|
||||
))
|
||||
continue
|
||||
|
||||
# 检查是否存在
|
||||
bucket = self.datasource._data.get(self.node_type, {})
|
||||
existing = bucket.get(node.data_id)
|
||||
if existing is None:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error=f"Item {node.data_id} not found"
|
||||
))
|
||||
continue
|
||||
|
||||
# 合并更新
|
||||
existing.update(data_dict)
|
||||
existing["id"] = node.data_id
|
||||
existing["_id"] = node.data_id
|
||||
if "type" not in existing:
|
||||
existing["type"] = self.node_type
|
||||
|
||||
# 更新缓存
|
||||
bucket[node.data_id] = existing
|
||||
self.datasource._dirty_types.add(self.node_type)
|
||||
|
||||
results.append(TaskResult.success(
|
||||
node_id=node.node_id,
|
||||
data_id=node.data_id
|
||||
))
|
||||
|
||||
except Exception as e:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error=f"Update failed: {str(e)}"
|
||||
))
|
||||
|
||||
# 自动保存到文件
|
||||
await self._save_to_file()
|
||||
|
||||
return results
|
||||
|
||||
async def _save_to_file(self) -> None:
|
||||
"""
|
||||
将缓存数据写回 JSONL 文件
|
||||
|
||||
在批量操作（create_all/update_all/delete_all）后自动调用。
|
||||
"""
|
||||
if self.datasource.read_only:
|
||||
return
|
||||
|
||||
if self.node_type not in self.datasource._dirty_types:
|
||||
return
|
||||
|
||||
items = self.datasource._data.get(self.node_type, {})
|
||||
if not items:
|
||||
# 如果缓存为空，可能是删除了所有数据，仍需写入空文件
|
||||
pass
|
||||
|
||||
try:
|
||||
# 写入到 {node_type}_1.jsonl
|
||||
file_path = self.datasource.dir_path / f"{self.node_type}_1.jsonl"
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(file_path, "w", encoding="utf-8") as f:
|
||||
for item in items.values():
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
|
||||
# 清除该类型的 dirty 标记
|
||||
self.datasource._dirty_types.discard(self.node_type)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error saving {self.node_type}: {str(e)}")
|
||||
# 不抛出异常，继续执行
|
||||
|
||||
async def delete_all(
|
||||
self,
|
||||
nodes: List['SyncNode']
|
||||
) -> List[TaskResult]:
|
||||
"""
|
||||
批量删除节点
|
||||
|
||||
Args:
|
||||
nodes: 需要删除的节点列表（DataSource 已筛选为 DELETE + PENDING）
|
||||
|
||||
Returns:
|
||||
List[TaskResult]，每个结果关联 node_id
|
||||
"""
|
||||
results = []
|
||||
|
||||
for node in nodes:
|
||||
try:
|
||||
if not node.data_id:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error="Cannot delete node without data_id"
|
||||
))
|
||||
continue
|
||||
|
||||
# 删除缓存
|
||||
bucket = self.datasource._data.get(self.node_type, {})
|
||||
if node.data_id in bucket:
|
||||
del bucket[node.data_id]
|
||||
self.datasource._dirty_types.add(self.node_type)
|
||||
results.append(TaskResult.success(
|
||||
node_id=node.node_id
|
||||
))
|
||||
else:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error=f"Item {node.data_id} not found"
|
||||
))
|
||||
|
||||
except Exception as e:
|
||||
results.append(TaskResult.failed(
|
||||
node_id=node.node_id,
|
||||
error=f"Delete failed: {str(e)}"
|
||||
))
|
||||
|
||||
# 自动保存到文件
|
||||
await self._save_to_file()
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user