ecm_sync_system/sync_state_machine/datasource/jsonl/datasource.py

"""
JsonlDataSource - JSONL 文件数据源

职责：
- 提供 JSONL 文件目录管理
- 实现 DataSource 接口（load_all, sync_all）
- 管理内存缓存和文件写入
- 统一异常处理和状态管理
"""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Dict, Any, Set, Callable, List, Optional

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()
        self._node_items_cache: Dict[str, Dict[str, Any]] = {}

    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}"
        )

    def invalidate_node_cache(self, node_type: str) -> None:
        self._node_items_cache.pop(node_type, None)

    def _matching_files(self, node_type: str) -> List[Path]:
        if not self.dir_path.exists():
            return []
        pattern = re.compile(rf"^{re.escape(node_type)}_\d+\.jsonl$")
        return sorted(
            file_path
            for file_path in self.dir_path.iterdir()
            if pattern.match(file_path.name)
        )

    def _build_cache_entry(
        self,
        *,
        node_type: str,
        extract_id: Callable[[Dict[str, Any]], str],
    ) -> Dict[str, Any]:
        files = self._matching_files(node_type)
        signature = tuple(
            (file_path.name, file_path.stat().st_mtime_ns, file_path.stat().st_size)
            for file_path in files
        )
        cached = self._node_items_cache.get(node_type)
        if cached and cached.get("signature") == signature:
            return cached

        items: List[Dict[str, Any]] = []
        data_id_index: Dict[str, List[Dict[str, Any]]] = {}
        project_index: Dict[str, List[Dict[str, Any]]] = {}
        no_project_items: List[Dict[str, Any]] = []

        for file_path in files:
            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)
                    except json.JSONDecodeError:
                        print(f"Warning: Invalid JSON in {file_path.name}:{line_num}")
                        continue

                    items.append(item)

                    item_id = extract_id(item)
                    if item_id:
                        data_id_index.setdefault(str(item_id), []).append(item)

                    project_id = item.get("project_id")
                    if project_id:
                        project_index.setdefault(str(project_id), []).append(item)
                    else:
                        no_project_items.append(item)

        entry = {
            "signature": signature,
            "items": items,
            "data_id_index": data_id_index,
            "project_index": project_index,
            "no_project_items": no_project_items,
        }
        self._node_items_cache[node_type] = entry
        return entry

    def get_items_for_load(
        self,
        *,
        node_type: str,
        extract_id: Callable[[Dict[str, Any]], str],
        data_id_filter: Optional[List[str]] = None,
        filter_on_project_id: bool = False,
    ) -> List[Dict[str, Any]]:
        entry = self._build_cache_entry(node_type=node_type, extract_id=extract_id)
        filters = [str(item) for item in (data_id_filter or []) if str(item)]
        if not filters:
            return list(entry["items"])

        if not filter_on_project_id:
            matched: List[Dict[str, Any]] = []
            for data_id in filters:
                matched.extend(entry["data_id_index"].get(data_id, []))
            return matched

        matched = list(entry["no_project_items"])
        for project_id in filters:
            matched.extend(entry["project_index"].get(project_id, []))
        return matched
    
    # load_all 和 sync_all 继承自 BaseDataSource
    # 保存逻辑在 BaseJsonlHandler 中实现