feat: 添加 FileTracker — 基于 SHA256 的文件变更追踪
This commit is contained in:
@@ -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()
|
||||||
Reference in New Issue
Block a user