feat: 添加 FileTracker — 基于 SHA256 的文件变更追踪

2026-07-11 19:39:12 +08:00
parent 78ad21d311
commit 3e2c61402a
2 changed files with 221 additions and 0 deletions
@@ -0,0 +1,117 @@
+"""文件变更追踪 — 基于 SHA256 + mtime 判断文件是否需要重新入库."""
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from threading import Lock
+from typing import TypedDict
+
+logger = logging.getLogger("md-vector-db")
+
+
+class FileRecord(TypedDict):
+    """追踪记录."""
+
+    hash: str
+    mtime: float
+    size: int
+    ingested_at: str
+
+
+class FileTracker:
+    """文件入库追踪器 — JSON 文件持久化.
+
+    通过比对文件内容的 SHA256 哈希判断文件是否变更。
+    使用 threading.Lock 保护 JSON 文件的并发读写。
+    """
+
+    def __init__(self, db_path: str = "./ingest_tracker.json"):
+        self._db_path = Path(db_path)
+        self._lock = Lock()
+        self._records: dict[str, FileRecord] = {}
+        self._load()
+
+    # -- 公开 API --
+
+    def is_stale(self, file_path: str) -> bool:
+        """检查文件是否需要重新入库.
+
+        Returns:
+            True: 文件不存在 / 无记录 / 内容已变更
+            False: 文件未变更且记录存在
+        """
+        path = Path(file_path)
+        if not path.exists():
+            return True
+        record = self.get_record(file_path)
+        if record is None:
+            return True
+        current_hash = self.compute_hash(file_path)
+        return current_hash != record["hash"]
+
+    def get_record(self, file_path: str) -> FileRecord | None:
+        """获取文件的追踪记录（无记录返回 None）."""
+        return self._records.get(self._abs_key(file_path))
+
+    def mark_ingested(self, file_path: str) -> None:
+        """标记文件已入库（创建或更新追踪记录）."""
+        path = Path(file_path)
+        if not path.exists():
+            logger.warning("标记已入库时文件不存在: %s", file_path)
+            return
+        stat = path.stat()
+        record: FileRecord = {
+            "hash": self.compute_hash(file_path),
+            "mtime": stat.st_mtime,
+            "size": stat.st_size,
+            "ingested_at": self._now_iso(),
+        }
+        with self._lock:
+            self._records[self._abs_key(file_path)] = record
+            self._save()
+
+    def remove_record(self, file_path: str) -> None:
+        """移除文件的追踪记录."""
+        key = self._abs_key(file_path)
+        with self._lock:
+            if key in self._records:
+                del self._records[key]
+                self._save()
+
+    @staticmethod
+    def compute_hash(file_path: str) -> str:
+        """计算文件 SHA256 哈希（分块读取，适合大文件）."""
+        sha = hashlib.sha256()
+        with open(file_path, "rb") as f:
+            while chunk := f.read(8192):
+                sha.update(chunk)
+        return sha.hexdigest()
+
+    # -- 内部 --
+
+    def _abs_key(self, file_path: str) -> str:
+        """生成标准化 key — 使用绝对路径."""
+        return str(Path(file_path).resolve())
+
+    @staticmethod
+    def _now_iso() -> str:
+        return datetime.now(timezone.utc).isoformat()
+
+    def _load(self) -> None:
+        if self._db_path.exists():
+            try:
+                with open(self._db_path, "r", encoding="utf-8") as f:
+                    self._records = json.load(f)
+            except (json.JSONDecodeError, OSError) as e:
+                logger.warning("tracker 文件损坏，重置为空: %s", e)
+                self._records = {}
+
+    def _save(self) -> None:
+        try:
+            with open(self._db_path, "w", encoding="utf-8") as f:
+                json.dump(self._records, f, ensure_ascii=False, indent=2)
+        except OSError as e:
+            logger.error("无法写入 tracker 文件: %s", e)
@@ -0,0 +1,104 @@
+"""文件变更追踪器测试."""
+import hashlib
+from pathlib import Path
+
+from src.core.file_tracker import FileTracker
+
+
+def test_compute_file_hash(tmp_path: Path):
+    """计算文件 SHA256 哈希."""
+    file = tmp_path / "test.md"
+    file.write_text("hello world", encoding="utf-8")
+    result = FileTracker.compute_hash(str(file))
+    expected = hashlib.sha256(b"hello world").hexdigest()
+    assert result == expected
+
+
+def test_compute_hash_deterministic(tmp_path: Path):
+    """同一文件内容产生相同哈希."""
+    file = tmp_path / "a.md"
+    file.write_text("same content", encoding="utf-8")
+    assert FileTracker.compute_hash(str(file)) == FileTracker.compute_hash(str(file))
+
+
+def test_hash_changes_with_content(tmp_path: Path):
+    """内容变更导致哈希不同."""
+    file = tmp_path / "b.md"
+    file.write_text("v1", encoding="utf-8")
+    h1 = FileTracker.compute_hash(str(file))
+    file.write_text("v2", encoding="utf-8")
+    h2 = FileTracker.compute_hash(str(file))
+    assert h1 != h2
+
+
+def test_is_stale_new_file(tmp_path: Path):
+    """新文件（无记录）视为过期."""
+    file = tmp_path / "new.md"
+    file.write_text("content", encoding="utf-8")
+    tracker = FileTracker(tmp_path / "tracker.json")
+    record = tracker.get_record(str(file))
+    assert record is None
+    assert tracker.is_stale(str(file)) is True
+
+
+def test_is_stale_unchanged_file(tmp_path: Path):
+    """未修改文件视为未过期."""
+    file = tmp_path / "unchanged.md"
+    file.write_text("stable", encoding="utf-8")
+    tracker = FileTracker(tmp_path / "tracker.json")
+    tracker.mark_ingested(str(file))
+    assert tracker.is_stale(str(file)) is False
+
+
+def test_is_stale_modified_file(tmp_path: Path):
+    """修改后文件视为过期."""
+    file = tmp_path / "modified.md"
+    file.write_text("v1", encoding="utf-8")
+    tracker = FileTracker(tmp_path / "tracker.json")
+    tracker.mark_ingested(str(file))
+    file.write_text("v2", encoding="utf-8")
+    assert tracker.is_stale(str(file)) is True
+
+
+def test_mark_ingested_updates_record(tmp_path: Path):
+    """mark_ingested 创建/更新记录."""
+    file = tmp_path / "x.md"
+    file.write_text("hello", encoding="utf-8")
+    tracker = FileTracker(tmp_path / "tracker.json")
+    tracker.mark_ingested(str(file))
+    record = tracker.get_record(str(file))
+    assert record is not None
+    assert record["hash"] == FileTracker.compute_hash(str(file))
+    assert "ingested_at" in record
+
+
+def test_persistence_across_instances(tmp_path: Path):
+    """tracker 数据持久化到 JSON，跨实例可读."""
+    file = tmp_path / "p.md"
+    file.write_text("persist me", encoding="utf-8")
+    db_path = tmp_path / "tracker.json"
+
+    t1 = FileTracker(db_path)
+    t1.mark_ingested(str(file))
+
+    t2 = FileTracker(db_path)
+    assert t2.is_stale(str(file)) is False
+
+
+def test_file_deleted_considered_stale(tmp_path: Path):
+    """文件被删除后视为过期."""
+    file = tmp_path / "tmp.md"
+    file.write_text("temp", encoding="utf-8")
+    tracker = FileTracker(tmp_path / "tracker.json")
+    tracker.mark_ingested(str(file))
+    file.unlink()
+    assert tracker.is_stale(str(file)) is True
+
+
+def test_binary_file_hash(tmp_path: Path):
+    """二进制文件也能正确计算哈希（如 PDF）."""
+    file = tmp_path / "doc.pdf"
+    file.write_bytes(b"\x00\x01\x02\x03")
+    h = FileTracker.compute_hash(str(file))
+    assert len(h) == 64
+    assert h == hashlib.sha256(b"\x00\x01\x02\x03").hexdigest()