From 3e2c61402a0e447f0fffde1b08838d869fe95889 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com> Date: Sat, 11 Jul 2026 19:39:12 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20FileTracker=20?= =?UTF-8?q?=E2=80=94=20=E5=9F=BA=E4=BA=8E=20SHA256=20=E7=9A=84=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E5=8F=98=E6=9B=B4=E8=BF=BD=E8=B8=AA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/file_tracker.py | 117 +++++++++++++++++++++++++++++++++++++ tests/test_file_tracker.py | 104 +++++++++++++++++++++++++++++++++ 2 files changed, 221 insertions(+) create mode 100644 src/core/file_tracker.py create mode 100644 tests/test_file_tracker.py diff --git a/src/core/file_tracker.py b/src/core/file_tracker.py new file mode 100644 index 0000000..637a82c --- /dev/null +++ b/src/core/file_tracker.py @@ -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) diff --git a/tests/test_file_tracker.py b/tests/test_file_tracker.py new file mode 100644 index 0000000..58cee72 --- /dev/null +++ b/tests/test_file_tracker.py @@ -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()