""" BaseJsonlHandler - JSONL Handler 基类 职责： - 查找和读取 JSONL 文件 - 实现批量 Handler 接口（load, create_all, update_all, delete_all） - 使用 datasource 的内存缓存 - 返回 List[TaskResult]（统一状态管理） - 捕获所有异常并转换为 TaskResult.failed() """ from __future__ import annotations import json 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_class_or_node_type, node_class: type | None = None, schema: type | None = None, node_type: str | None = None, ): if isinstance(node_class_or_node_type, str): resolved_node_type = node_class_or_node_type resolved_node_class = node_class resolved_schema = schema else: resolved_node_class = node_class_or_node_type resolved_node_type = node_type or getattr(resolved_node_class, "node_type", None) resolved_schema = schema or getattr(resolved_node_class, "schema", None) super().__init__(resolved_node_type, resolved_node_class, resolved_schema) self.datasource = datasource 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 data_id_filter = self.get_data_id_filter() items = self.datasource.get_items_for_load( node_type=self.node_type, extract_id=self.extract_id, data_id_filter=data_id_filter, filter_on_project_id=self.node_type != "project", ) for item in items: item_id = self.extract_id(item) if self.should_skip_by_data_id_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 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) self.datasource.invalidate_node_cache(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