strepsiades
344 lines
11 KiB
Python
344 lines
11 KiB
Python
"""
|
||
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
|
||
|
||
def _should_skip_item_by_project_filter(self, item: Dict[str, Any], item_id: str) -> bool:
|
||
data_id_filter = set(self.get_data_id_filter())
|
||
if not data_id_filter:
|
||
return False
|
||
if self.node_type == "project":
|
||
return item_id not in data_id_filter
|
||
|
||
project_id = item.get("project_id")
|
||
if project_id:
|
||
return str(project_id) not in data_id_filter
|
||
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
|